gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

disabling interleaving

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-13 17:16:00 -07:00
parent d487cf31ce
commit 21a9b92c87
4 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/python/tensorflow/dot.cpp
View File

@@ -23,7 +23,7 @@ const char* src =
R"(
const tunable int32 TM = {64, 128};
const tunable int32 TN = {64, 128};
const tunable int32 TK = {16};
const tunable int32 TK = {32};
const tunable int32 GZ = {1};
void matmul(restrict read_only fp16 *A, restrict read_only fp16 *B,
@@ -117,10 +117,10 @@ class BlockSparseGemmOp : public OpKernel {
return 2.*M*N*K / ts * 1e-3;
};
// just-in-time compile source-code
jit.autotune("matmul", src, benchmark);
// jit.autotune("matmul", src, benchmark);
// jit.add_module("matmul", src, {4, 2, 8, 4, 2, 32, 1, 4, 1, 1, 8, 8, 8, 1});
// jit.add_module("matmul", src, {32, 2, 128, 32, 2, 128, 2, 2, 2, 2, 4, 8, 4, 1});
jit.add_module("matmul", src, {16, 4, 128, 16, 4, 128, 2, 2, 2, 2, 8, 32, 8, 1});
// jit.add_module("matmul", src, {8, 8, 128, 16, 8, 128, 2, 2, 2, 2, 16, 32, 8, 1 });
triton::driver::kernel* kernel = jit.get_function("matmul");
triton::jit::launch_information info = jit.get_launch_info("matmul");
std::cout << benchmark(kernel, info) << std::endl;;

11
lib/codegen/selection.cpp
View File

@@ -501,8 +501,8 @@ void selection::init_axes(ir::value *v, IRBuilder<> &builder, Value *u_thread_id
unsigned num_rep_0 = shapes[0]->get_value() / hmma_bts_0;
unsigned num_rep_1 = shapes[1]->get_value() / hmma_bts_1;
// size of each pack (interleaving)
pack_size_0_ = std::min<unsigned>(num_rep_0, 2);
pack_size_1_ = std::min<unsigned>(num_rep_1, 2);
pack_size_0_ = std::min<unsigned>(num_rep_0, 1);
pack_size_1_ = std::min<unsigned>(num_rep_1, 1);
// number of packs (interleaving)
num_packs_0_ = num_rep_0 / pack_size_0_;
num_packs_1_ = num_rep_1 / pack_size_1_;
@@ -531,9 +531,14 @@ void selection::init_axes(ir::value *v, IRBuilder<> &builder, Value *u_thread_id
// a offset
offset_a_i_ = builder.CreateAdd(warp_offset_i, builder.CreateAdd(pair_a_off, in_pair_off_a));
offset_a_k_ = builder.CreateAnd(u_thread_id, _3);
// b offsets
// // b offsets
offset_b_j_ = builder.CreateAdd(warp_offset_j, builder.CreateAdd(pair_b_off, in_pair_off_b));
offset_b_k_ = builder.CreateAnd(u_thread_id, _3);
// offset_a_i_ = builder.getInt32(0);
// offset_a_k_ = builder.getInt32(0);
// offset_b_j_ = builder.getInt32(0);
// offset_b_k_ = builder.getInt32(0);
// c offsets
Value *offset_c_i = builder.CreateAdd(builder.CreateAnd(u_thread_id, _1), offset_a_i_);
Value *offset_c_j = builder.CreateAdd(builder.CreateAnd(u_thread_id, _2),

7
lib/codegen/shmem_allocation.cpp
View File

@@ -23,10 +23,11 @@ unsigned shmem_allocation::is_ld_padded(ir::value *x) {
for(ir::user* user: x->get_users())
if(dynamic_cast<ir::dot_inst*>(user))
if(params_->get_fragment(user, 0) == tune::HMMA_FRAGMENT_C){
if(x == user->get_operand(0))
return 4;
if(x == user->get_operand(0)){
return 16;
}
else
return 4;
return 16;
}
return 0;
}

2
lib/driver/module.cpp
View File

@@ -255,7 +255,7 @@ std::string cu_module::compile_llvm_module(llvm::Module* module) {
cu_module::cu_module(driver::context * context, llvm::Module* ll_module): cu_module(context, compile_llvm_module(ll_module)) { }
cu_module::cu_module(driver::context * context, std::string const & source) : module(context, CUmodule(), true), source_(source){
// std::cout << source << std::endl;
std::cout << source << std::endl;
cu_context::context_switcher ctx_switch(*context);
// JIT compile source-code
CUjit_option opt[] = {CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES, CU_JIT_ERROR_LOG_BUFFER};