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

[codegen] bugfix in builder insert point for predicated instructions

Browse Source
This commit is contained in:
Philippe Tillet
2019-05-04 12:09:27 -04:00
parent 30833c18f1
commit 4813bb007c
3 changed files with 8 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
examples/cpp/dot.cpp
View File

@@ -53,26 +53,11 @@ void matmul(restrict read_only fp32 *A, restrict read_only fp32 *B, fp32 *C,
fp32 b[TN, 1] = checkb ? *pb : 0; fp32 b[TN, 1] = checkb ? *pb : 0;
c = dot(a, trans(b), c); c = dot(a, trans(b), c);
} }
int32 ridx = get_range_id(0);
int32 ridy = get_range_id(1);
fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis]; fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
int32 *plock = locks + ridx + ridy*grid0;
while(__atomic_cas(plock, 0, 1));
int32 *pcount = plock + grid0*grid1;
int32 count = *pcount;
int32 countp1 = select(count == GZ - 1, 0, count + 1);
int1 checkc0[TM] = rxc < M; int1 checkc0[TM] = rxc < M;
int1 checkc1[TN] = ryc < N; int1 checkc1[TN] = ryc < N;
int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :]; int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
if(count == 0) { @checkc *pc = c;
@checkc *pc = c;
*pcount = countp1;
}
else {
@checkc *pc = c + *pc;
*pcount = countp1;
}
__atomic_cas(plock, 1, 0);
} }
)"; )";

1
include/triton/jit.h
View File

@@ -58,6 +58,7 @@ public:
target_(target) { } target_(target) { }
void target_independent(ir::module &module) { void target_independent(ir::module &module) {
// ir::print(module, std::cout);
optimize_dot.run(module); optimize_dot.run(module);
optimize_trans.run(module); optimize_trans.run(module);
} }

8
lib/codegen/selection.cpp
View File

@@ -632,11 +632,13 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
LLVMContext &ctx = builder.getContext(); LLVMContext &ctx = builder.getContext();
Function *fn = block->getParent(); Function *fn = block->getParent();
ir::value *mask = ins->get_mask_pred(); ir::value *mask = ins->get_mask_pred();
BasicBlock *last_block = nullptr;
auto set_mask_insert_pt = [&](indices_t idx){ auto set_mask_insert_pt = [&](indices_t idx){
if(mask){ if(mask){
distributed_tile *mask_tile = (distributed_tile*)tmap_.at(ins->get_mask_pred()); distributed_tile *mask_tile = (distributed_tile*)tmap_.at(ins->get_mask_pred());
BasicBlock *block = pmap_.at({mask_tile, idx}); BasicBlock *block = pmap_.at({mask_tile, idx});
builder.SetInsertPoint(block->getTerminator()); builder.SetInsertPoint(block->getTerminator());
last_block = last_block_.at({mask_tile, idx});
} }
}; };
// store // store
@@ -646,7 +648,6 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
ptr->for_each([&](indices_t idx){ ptr->for_each([&](indices_t idx){
set_mask_insert_pt(idx); set_mask_insert_pt(idx);
StoreInst *store = new StoreInst(value->get_value(idx), ptr->get_value(idx)); StoreInst *store = new StoreInst(value->get_value(idx), ptr->get_value(idx));
// store->setAlignment(16);
builder.Insert(store); builder.Insert(store);
}); });
} }
@@ -847,8 +848,11 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
}); });
} }
} }
if(mask) if(mask){
builder.SetInsertPoint(block); builder.SetInsertPoint(block);
if(last_block)
builder.SetInsertPoint(last_block);
}
} }
void selection::lower_instruction(ir::instruction *src, IRBuilder<> &builder) { void selection::lower_instruction(ir::instruction *src, IRBuilder<> &builder) {