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

[CODEGEN] Fixed bug that caused conditional operator to not always

Browse Source 
properly mask load operations

Also includes minor improvement to benchmarking infrastructure
This commit is contained in:
Philippe Tillet
2021-03-07 14:53:48 -05:00
parent d1d09566b1
commit 5b9afaa688
9 changed files with 146 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
lib/codegen/selection/generator.cc
View File

@@ -1616,20 +1616,22 @@ void generator::visit_make_range(ir::make_range* x) {
}
}
void generator::visit_undef_value(ir::undef_value *ud) {
vals_[ud][{}] = llvm::UndefValue::get(cvt(ud->get_type()));
void generator::visit_undef_value(ir::undef_value *x) {
Type* ty = cvt(x->get_type()->get_scalar_ty());
for(indices_t idx: idxs_.at(x))
vals_[x][idx] = llvm::UndefValue::get(ty);
}
void generator::visit_constant_int(ir::constant_int *cst){
Type *ty = cvt(cst->get_type()->get_scalar_ty());
vals_[cst][{}] = ConstantInt::get(ty, cst->get_value());
void generator::visit_constant_int(ir::constant_int *x){
Type *ty = cvt(x->get_type()->get_scalar_ty());
for(indices_t idx: idxs_.at(x))
vals_[x][idx] = ConstantInt::get(ty, x->get_value());
}
void generator::visit_constant_fp(ir::constant_fp *cst){
Type *ty = cvt(cst->get_type()->get_scalar_ty());
vals_[cst][{}] = ConstantFP::get(ty, cst->get_value());
void generator::visit_constant_fp(ir::constant_fp *x){
Type *ty = cvt(x->get_type()->get_scalar_ty());
for(indices_t idx: idxs_.at(x))
vals_[x][idx] = ConstantFP::get(ty, x->get_value());
}
void generator::visit_alloc_const(ir::alloc_const *alloc) {

17
lib/codegen/transform/peephole.cc
View File

@@ -193,6 +193,22 @@ bool peephole::rewrite_gep_ptr_min_off_plus_off(ir::instruction *value, ir::buil
return false;
}
bool peephole::rewrite_select_masked_load(ir::instruction *value, ir::builder& builder){
auto select = dynamic_cast<ir::select_inst*>(value);
if(!select)
return false;
auto if_value = dynamic_cast<ir::masked_load_inst*>(select->get_if_value_op());
if(!if_value)
return false;
if(select->get_pred_op() != if_value->get_mask_operand())
return false;
builder.set_insert_point(select);
ir::value* new_load = builder.create_masked_load(if_value->get_pointer_operand(),
if_value->get_mask_operand(),
select->get_else_value_op());
select->replace_all_uses_with(new_load);
return true;
}
void peephole::run(ir::module &mod) {
ir::builder &builder = mod.get_builder();
@@ -230,6 +246,7 @@ void peephole::run(ir::module &mod) {
// was_modified = was_modified || rewrite_trans_phi(i, builder);
was_modified = was_modified || rewrite_unit_red(i, builder);
was_modified = was_modified || rewrite_gep_ptr_min_off_plus_off(i, builder);
was_modified = was_modified || rewrite_select_masked_load(i, builder);
if(tgt_->as_nvidia()->sm() >= 80)
was_modified = was_modified || rewrite_load_to_shared(i, builder);
if(was_modified)

55
lib/lang/code_gen.cc
View File

@@ -267,23 +267,56 @@ void Generator::VisitTransOp(TransOp *trans) {
}
void Generator::VisitConditionalOp(ConditionalOp* condOp) {
// auto &instructions = bld_->get_insert_block()->get_inst_list();
auto &instructions = bld_->get_insert_block()->get_inst_list();
VisitExpr(condOp->cond_);
ir::value* cond = ret_;
ir::value* true_cond = ret_;
ir::instruction* start = instructions.back();
VisitExpr(condOp->exprTrue_);
ir::value* true_val = ret_;
VisitExpr(condOp->exprFalse_);
ir::value* false_val = ret_;
if(ir::unmasked_load_inst* ld = dynamic_cast<ir::unmasked_load_inst*>(true_val)) {
if(true_val->get_type()->is_tile_ty() && !false_val->get_type()->is_tile_ty())
false_val = bld_->create_splat(false_val, cond->get_type()->get_tile_shapes());
ir::value* new_ld = bld_->create_masked_load(ld->get_pointer_operand(), cond, false_val);
ld->replace_all_uses_with(new_ld);
ld->erase_from_parent();
return set_ret(new_ld);
auto begin = std::find(instructions.begin(), instructions.end(), start);
bool is_in_true_cond = true;
for(auto it = begin; it != instructions.end(); it++){
ir::instruction* instr = *it;
// we mask load with `cond` when used to compute true_value
// we mask load with `!cond` when used to compute false_value
if(auto ld = dynamic_cast<ir::unmasked_load_inst*>(instr)){
bld_->set_insert_point(ld);
ir::type* ty = ld->get_type();
ir::value* cond = is_in_true_cond ? true_cond : true_cond;
ir::value* ptr = ld->get_pointer_operand();
ir::value* else_val = ir::undef_value::get(ty);
ir::value* masked_ld = bld_->create_masked_load(ptr, cond, else_val);
ld->replace_all_uses_with(masked_ld);
ld->erase_from_parent();
if(true_val == ld)
true_val = masked_ld;
if(false_val == ld)
false_val = masked_ld;
it = std::find(instructions.begin(), instructions.end(), masked_ld);
}
if(instr == true_val)
is_in_true_cond = false;
}
return set_ret(bld_->create_select(cond, true_val, false_val));
// return error_not_implemented();
bld_->set_insert_point(bld_->get_insert_block());
return set_ret(bld_->create_select(true_cond, true_val, false_val));
// VisitExpr(condOp->cond_);
// ir::value* cond = ret_;
// VisitExpr(condOp->exprTrue_);
// ir::value* true_val = ret_;
// VisitExpr(condOp->exprFalse_);
// ir::value* false_val = ret_;
// if(ir::unmasked_load_inst* ld = dynamic_cast<ir::unmasked_load_inst*>(true_val)) {
// if(true_val->get_type()->is_tile_ty() && !false_val->get_type()->is_tile_ty())
// false_val = bld_->create_splat(false_val, cond->get_type()->get_tile_shapes());
// ir::value* new_ld = bld_->create_masked_load(ld->get_pointer_operand(), cond, false_val);
// ld->replace_all_uses_with(new_ld);
// ld->erase_from_parent();
// return set_ret(new_ld);
// }
// return set_ret(bld_->create_select(cond, true_val, false_val));
}
void Generator::VisitFuncCall(FuncCall* funcCall) {

1
lib/runtime/function.cc
View File

@@ -156,6 +156,7 @@ std::tuple<std::shared_ptr<driver::module>,
layouts.run(ir);
peephole.run(ir);
dce.run(ir);
// ir::print(ir, std::cout);
if(target->is_gpu())
cts.run(ir);
align.run(ir);