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

[test] [reduce] added test for 1D reduction

Browse Source
This commit is contained in:
Philippe Tillet
2019-10-20 01:01:53 -04:00
parent 23db500edf
commit abe3fbb480
8 changed files with 79 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lib/codegen/analysis/axes.cc
View File

@@ -106,6 +106,8 @@ void axes::update_graph_elementwise(ir::instruction *i) {
}
void axes::update_graph_no_edge(ir::instruction *i) {
if(!i->get_type()->is_tile_ty())
return;
auto rank = i->get_type()->get_tile_rank();
for(unsigned d = 0; d < rank; d++)
graph_.add_edge({i, d}, {i, d});

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

@@ -825,13 +825,12 @@ void generator::visit_reduce_inst(ir::reduce_inst* x) {
}
tgt_->add_barrier(mod_, *builder_);
// write back
distributed_tile* x_tile = (distributed_tile*)tmap_.at(x);
x_tile->for_each([&](indices_t idx) {
for_each(x, [&](indices_t idx) {
indices_t red_idx = idx;
red_idx.insert(red_idx.begin() + axis, builder_->getInt32(0));
Value *read_offset = shared_tile::shared_offset(*builder_, stile->get_shapes(), stile->get_perm(), stile->get_order(), red_idx);
Value *read_ptr = builder_->CreateGEP(base_ptr, read_offset);
x_tile->set_value(idx, builder_->CreateLoad(read_ptr));
set_value(x, idx, builder_->CreateLoad(read_ptr));
});
}