Skip to content

Commit

Permalink
[amdgpu] Add convert addressspace pass related unit test (#7023)
Browse files Browse the repository at this point in the history 
Issue: ##6434

### Brief Summary
These unit tests are for #6486

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
@galeselee @pre-commit-ci
galeselee and pre-commit-ci[bot] authored Jan 6, 2023
1 parent 5d8aab9 commit 328cacf
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 5 deletions.
9 changes: 4 additions & 5 deletions taichi/runtime/llvm/llvm_context_pass.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/SourceMgr.h"

#if defined(TI_WITH_AMDGPU)
#include "taichi/rhi/amdgpu/amdgpu_context.h"
Expand All @@ -19,7 +20,7 @@ namespace lang {
using namespace llvm;
#if defined(TI_WITH_AMDGPU)
struct AMDGPUConvertAllocaInstAddressSpacePass : public FunctionPass {
static char ID;
static inline char ID{0};
AMDGPUConvertAllocaInstAddressSpacePass() : FunctionPass(ID) {
}
bool runOnFunction(llvm::Function &f) override {
Expand All @@ -38,7 +39,7 @@ struct AMDGPUConvertAllocaInstAddressSpacePass : public FunctionPass {
}
for (auto &allocainst : alloca_inst_vec) {
auto alloca_type = allocainst->getAllocatedType();
IRBuilder<> builder(allocainst);
llvm::IRBuilder<> builder(allocainst);
auto *new_alloca = builder.CreateAlloca(alloca_type, (unsigned)5);
auto new_type = llvm::PointerType::get(alloca_type, (unsigned)0);
new_alloca->setAlignment(Align(allocainst->getAlign().value()));
Expand All @@ -52,7 +53,7 @@ struct AMDGPUConvertAllocaInstAddressSpacePass : public FunctionPass {
};

struct AMDGPUConvertFuncParamAddressSpacePass : public ModulePass {
static char ID;
static inline char ID{0};
AMDGPUConvertFuncParamAddressSpacePass() : ModulePass(ID) {
}
bool runOnModule(llvm::Module &M) override {
Expand Down Expand Up @@ -127,8 +128,6 @@ struct AMDGPUConvertFuncParamAddressSpacePass : public ModulePass {
}
};

char AMDGPUConvertAllocaInstAddressSpacePass::ID = 0;
char AMDGPUConvertFuncParamAddressSpacePass::ID = 0;
#endif

} // namespace lang
Expand Down
100 changes: 100 additions & 0 deletions tests/cpp/backends/amdgpu_device_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,22 @@
#include "taichi/rhi/amdgpu/amdgpu_driver.h"
#include "taichi/rhi/amdgpu/amdgpu_context.h"
#include "taichi/rhi/amdgpu/amdgpu_device.h"
#include "taichi/runtime/llvm/llvm_context_pass.h"

#include <llvm/IR/Instructions.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/LegacyPassManager.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Module.h>
#include <llvm/Pass.h>
#include <llvm/IRReader/IRReader.h>

#include "tests/cpp/program/test_program.h"

namespace taichi {
namespace lang {

TEST(AMDGPU, CreateDeviceAndAlloc) {
std::unique_ptr<amdgpu::AmdgpuDevice> device =
std::make_unique<amdgpu::AmdgpuDevice>();
Expand Down Expand Up @@ -87,6 +99,94 @@ TEST(AMDGPU, CreateContextAndGetMemInfo) {
EXPECT_GE(free_size, 0);
}

TEST(AMDGPU, ConvertAllocaInstAddressSpacePass) {
const std::string program =
"define dso_local void @runtime_add(double* %0, double* %1, double* %2) "
"#4 "
"{\n"
"%4 = alloca double*, align 8\n"
"%5 = alloca double*, align 8\n"
"%6 = alloca double*, align 8\n"
"store double* %0, double** %4, align 8\n"
"store double* %1, double** %5, align 8\n"
"store double* %2, double** %6, align 8\n"
"%7 = load double*, double** %4, align 8\n"
"%8 = load double, double* %7, align 8\n"
"%9 = load double*, double** %5, align 8\n"
"%10 = load double, double* %9, align 8\n"
"%11 = fadd contract double %8, %10\n"
"%12 = load double*, double** %6, align 8\n"
"store double %11, double* %12, align 8\n"
"ret void\n"
"}\n";
llvm::LLVMContext llvm_context;
llvm::SMDiagnostic diagnostic_err;
std::unique_ptr<llvm::Module> llvm_module = llvm::parseIR(
llvm::MemoryBuffer::getMemBuffer(program)->getMemBufferRef(),
diagnostic_err, llvm_context);
llvm::legacy::FunctionPassManager function_pass_manager(llvm_module.get());
function_pass_manager.add(new AMDGPUConvertAllocaInstAddressSpacePass());
function_pass_manager.doInitialization();
for (auto func = llvm_module->begin(); func != llvm_module->end(); ++func) {
function_pass_manager.run(*func);
}
function_pass_manager.doFinalization();
auto func = llvm_module->getFunction("runtime_add");
for (auto &bb : *func) {
for (llvm::Instruction &inst : bb) {
auto alloca_inst = llvm::dyn_cast<AllocaInst>(&inst);
if (!alloca_inst)
continue;
EXPECT_EQ(alloca_inst->getAddressSpace(), 5);
}
int cast_num = 0;
for (llvm::Instruction &inst : bb) {
auto cast_inst = llvm::dyn_cast<AddrSpaceCastInst>(&inst);
if (!cast_inst)
continue;
cast_num++;
}
EXPECT_EQ(cast_num, 3);
}
}

TEST(AMDGPU, ConvertFuncParamAddressSpacePass) {
const std::string program =
"define dso_local void @runtime_add(double* %0, double* %1, double* %2) "
"#4 "
"{\n"
"%4 = alloca double*, align 8\n"
"%5 = alloca double*, align 8\n"
"%6 = alloca double*, align 8\n"
"store double* %0, double** %4, align 8\n"
"store double* %1, double** %5, align 8\n"
"store double* %2, double** %6, align 8\n"
"%7 = load double*, double** %4, align 8\n"
"%8 = load double, double* %7, align 8\n"
"%9 = load double*, double** %5, align 8\n"
"%10 = load double, double* %9, align 8\n"
"%11 = fadd contract double %8, %10\n"
"%12 = load double*, double** %6, align 8\n"
"store double %11, double* %12, align 8\n"
"ret void\n"
"}\n";
llvm::LLVMContext llvm_context;
llvm::SMDiagnostic diagnostic_err;
std::unique_ptr<llvm::Module> llvm_module = llvm::parseIR(
llvm::MemoryBuffer::getMemBuffer(program)->getMemBufferRef(),
diagnostic_err, llvm_context);
llvm::legacy::PassManager module_pass_manager;
module_pass_manager.add(new AMDGPUConvertFuncParamAddressSpacePass());
module_pass_manager.run(*llvm_module);
auto func = llvm_module->getFunction("runtime_add");
for (llvm::Function::arg_iterator I = func->arg_begin(), E = func->arg_end();
I != E; ++I) {
if (I->getType()->getTypeID() == llvm::Type::PointerTyID) {
EXPECT_EQ(I->getType()->getPointerAddressSpace(), 1);
}
}
}

} // namespace lang
} // namespace taichi
#endif

0 comments on commit 328cacf

Please sign in to comment.