From 1b9a7a8e972ea0392e7b83a2ff92fabc2e7e1f3f Mon Sep 17 00:00:00 2001 From: Philippe Tillet Date: Thu, 7 Feb 2019 17:03:19 -0500 Subject: [PATCH] [code generation] added basic shared copy/read --- examples/matrix.cpp | 4 +++ include/codegen/selection.h | 19 +++++++++---- lib/codegen/selection.cpp | 55 +++++++++++++++++++++++++++++++------ 3 files changed, 63 insertions(+), 15 deletions(-) diff --git a/examples/matrix.cpp b/examples/matrix.cpp index e5ea0b1ba..51847bfb2 100644 --- a/examples/matrix.cpp +++ b/examples/matrix.cpp @@ -34,7 +34,11 @@ void test(fp32 *a, fp32 *b, fp32 *c, int32 M, int32 N, int32 K){\ fp32* pb[32, 8] = b + ry[:, newaxis] + rkb[newaxis, :]*K;\ fp32* pc[32, 32] = c + rx[:, newaxis] + ry[newaxis, :]*M;\ for(k = K; k >= 0; k = k - 8){\ + fp32 a[32, 8] = *pa;\ + fp32 b[32, 8] = *pb;\ C = C + 1;\ + pa = pa + 8*M;\ + pb = pb + 8*K;\ }\ *pc = C;\ }\ diff --git a/include/codegen/selection.h b/include/codegen/selection.h index 8b39ee150..5aea5564f 100644 --- a/include/codegen/selection.h +++ b/include/codegen/selection.h @@ -37,16 +37,23 @@ public: virtual void set_value(indices_t idx, llvm::Value *v) = 0; virtual llvm::Value* get_value(indices_t idx) = 0; -private: +protected: llvm::Type *ty_; shapes_t shapes_; }; class shared_tile: public tile { +private: + llvm::Value* shared_offset(indices_t idx); + public: - using tile::tile; - void set_value(indices_t idx, llvm::Value *v) { } - llvm::Value* get_value(indices_t idx) { return nullptr; } + shared_tile(llvm::Type* ty, const shapes_t &shapes, llvm::Value* ptr, llvm::IRBuilder<> &builder); + void set_value(indices_t, llvm::Value *); + llvm::Value* get_value(indices_t idx); + +private: + llvm::Value *ptr_; + llvm::IRBuilder<> &builder_; }; class distributed_tile: public tile{ @@ -85,9 +92,9 @@ private: void create_grids(std::vector &grids, std::map &references, ir::function *fn); - void create_tile(ir::value *v, llvm::IRBuilder<> &builder, const std::map &references, std::set &seen); + void create_tile(ir::value *v, llvm::IRBuilder<> &builder, const std::map &references, std::set &seen, llvm::Value *sh_mem_ptr); void init_axes(ir::value *i, llvm::IRBuilder<> &builder, llvm::Value *u_thread_id, llvm::Value *u_warp_id); - void init_grids(ir::function *fn, llvm::IRBuilder<> &builder); + void init_grids(ir::function *fn, llvm::IRBuilder<> &builder, llvm::Value *sh_mem_ptr); // lowering void lower_instruction(ir::instruction *src, llvm::IRBuilder<> &builder); diff --git a/lib/codegen/selection.cpp b/lib/codegen/selection.cpp index bba6e8e43..473588363 100644 --- a/lib/codegen/selection.cpp +++ b/lib/codegen/selection.cpp @@ -1,5 +1,6 @@ #include "codegen/selection.h" #include "codegen/tune.h" +#include "codegen/allocation.h" #include "llvm/IR/Module.h" #include "llvm/IR/IRBuilder.h" #include "ir/context.h" @@ -53,6 +54,28 @@ void distributed_tile::for_each(std::function fn) { fn(idx.first); } +/* Shared Tile */ +Value* shared_tile::shared_offset(indices_t idx) { + Value *result = builder_.getInt32(0); + result = builder_.CreateAdd(result, idx[0]); + for(size_t i = 1; i < idx.size(); i++) + result = builder_.CreateAdd(result, builder_.CreateMul(idx[i], builder_.getInt32(shapes_[i-1]))); + return result; +} + +shared_tile::shared_tile(Type *ty, const shapes_t &shapes, Value *ptr, llvm::IRBuilder<> &builder): tile(ty, shapes), ptr_(ptr), builder_(builder) { + +} + +void shared_tile::set_value(indices_t idx, Value *value) { + Value *ptr = builder_.CreateGEP(ptr_, shared_offset(idx)); + builder_.CreateStore(value, ptr); +} + +Value* shared_tile::get_value(indices_t idx) { + Value *ptr = builder_.CreateGEP(ptr_, shared_offset(idx)); + return builder_.CreateLoad(ptr); +} /* convert ir::type to Type */ Type *selection::llvm_type(ir::type *ty, LLVMContext &ctx) { @@ -281,19 +304,21 @@ void selection::create_grids(std::vector &grids, void selection::create_tile(ir::value *v, IRBuilder<> &builder, const std::map& references, - std::set &seen) { + std::set &seen, Value *sh_mem_ptr) { if(!v->get_type()->is_tile_ty() || !seen.insert(v).second) return; if(auto *user = dynamic_cast(v)) for(ir::value *op: user->ops()) - create_tile(op, builder, references, seen); + create_tile(op, builder, references, seen, sh_mem_ptr); LLVMContext &ctx = builder.getContext(); - bool is_shared = dynamic_cast(v); const auto& shapes = v->get_type()->get_tile_shapes(); Type* ty = llvm_type(v->get_type()->get_scalar_ty(), ctx); // create shared tile + bool is_shared = dynamic_cast(v); if(is_shared){ - tmap_.insert({v, new shared_tile(ty, shapes)}); + size_t offset = alloc_->get_offset(v); + Value *ptr = builder.CreateGEP(sh_mem_ptr, builder.getInt32(offset)); + tmap_.insert({v, new shared_tile(ty, shapes, ptr, builder)}); } // create distributed tile else { @@ -318,7 +343,7 @@ void selection::create_tile(ir::value *v, IRBuilder<> &builder, } } -void selection::init_grids(ir::function *fn, IRBuilder<> &builder){ +void selection::init_grids(ir::function *fn, IRBuilder<> &builder, Value *sh_mem_ptr){ // fetch linear ID Module *mod = builder.GetInsertBlock()->getParent()->getParent(); Function *get_thread_id = Intrinsic::getDeclaration(mod, Intrinsic::nvvm_read_ptx_sreg_tid_x); @@ -338,7 +363,7 @@ void selection::init_grids(ir::function *fn, IRBuilder<> &builder){ for(ir::instruction *i: block->get_inst_list()){ if(!i->get_type()->is_tile_ty()) continue; - create_tile(i, builder, references, seen); + create_tile(i, builder, references, seen, sh_mem_ptr); } } @@ -411,7 +436,10 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> & } // copy to shared else if(dynamic_cast(ins)) { - + distributed_tile* in = (distributed_tile*)tmap_.at(ins->get_operand(0)); + in->for_each([&](indices_t idx){ + ti->set_value(idx, in->get_value(idx)); + }); } // element-wise else { @@ -444,6 +472,7 @@ void selection::run(ir::module &src, Module &dst){ vmap_.clear(); LLVMContext &dst_ctx = dst.getContext(); IRBuilder<> dst_builder(dst_ctx); + // iterate over functions for(ir::function *fn: src.get_function_list()) { // create LLVM function @@ -457,9 +486,17 @@ void selection::run(ir::module &src, Module &dst){ BasicBlock *dst_block = BasicBlock::Create(dst_ctx, block->get_name(), dst_fn); vmap_[block] = dst_block; } - // create grids dst_builder.SetInsertPoint((BasicBlock*)vmap_[fn->blocks()[0]]); - init_grids(fn, dst_builder); + // allocate shared memory + Type *int_8_ty = Type::getInt8Ty(dst_ctx); + ArrayType *array_ty = ArrayType::get(int_8_ty, alloc_->get_allocated_size()); + Type *ptr_ty = PointerType::get(int_8_ty, 3); + GlobalVariable *sh_mem_array = + new GlobalVariable(*dst_fn->getParent(), array_ty, false, GlobalVariable::InternalLinkage, + nullptr, "__shared_ptr", nullptr, GlobalVariable::NotThreadLocal, 3); + Value *sh_mem_ptr = dst_builder.CreateBitCast(sh_mem_array, ptr_ty); + // create grids + init_grids(fn, dst_builder, sh_mem_ptr); // iterate through block for(ir::basic_block *block: fn->blocks()) { dst_builder.SetInsertPoint((BasicBlock*)vmap_[block]);