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

[codegen] more cleaning

Browse Source
This commit is contained in:
Philippe Tillet
2019-10-10 15:52:03 -04:00
parent a3f76b6eb1
commit 4efd0a3c6b
9 changed files with 148 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
include/triton/codegen/analysis/layout.h
View File

@@ -27,14 +27,35 @@ enum layout_type_t {
}; };
struct layout_t { struct layout_t {
layout_t(layout_type_t _type,
const std::vector<int>& _axes,
const std::vector<unsigned> &_shapes,
const std::vector<ir::value *> &values,
analysis::align* align);
layout_type_t type; layout_type_t type;
std::vector<int> axes; std::vector<int> axes;
std::vector<unsigned> shapes; std::vector<unsigned> shapes;
std::vector<int> order; std::vector<int> order;
std::map<int, int> mts; std::vector<int> mts;
std::map<int, int> nts; std::vector<int> nts;
std::map<int, int> fpw; std::vector<int> fpw;
std::map<int, int> wpt; std::vector<int> wpt;
};
struct layout_hmma_884_t: public layout_t {
layout_hmma_884_t(size_t num_warps,
const std::vector<int>& _axes,
const std::vector<unsigned>& _shapes,
const std::vector<ir::value *> &values,
analysis::align* align);
};
struct layout_scanline_t: public layout_t {
layout_scanline_t(size_t num_warps,
const std::vector<int>& _axes,
const std::vector<unsigned>& _shapes,
const std::vector<ir::value *> &values,
analysis::align* align);
}; };
class layout { class layout {
@@ -52,12 +73,13 @@ private:
public: public:
// constructor // constructor
layout(analysis::axes *axes, analysis::align *align, size_t num_warps); layout(analysis::axes *axes, analysis::align *align, size_t num_warps);
// accessors // accessors
unsigned layout_of(ir::value *value) const; unsigned layout_of(ir::value *value) const;
const std::vector<ir::value*>& values_of(unsigned id) const; const std::vector<ir::value*>& values_of(unsigned id) const;
size_t num_layouts() const; size_t num_layouts() const;
const layout_t& get(ir::value *v) const; const layout_t* get(ir::value *v) const;
std::map<size_t, layout_t> &get_all(); std::map<size_t, layout_t*> &get_all();
// execution // execution
void run(ir::module &mod); void run(ir::module &mod);
@@ -69,7 +91,7 @@ private:
tools::graph<ir::value*> graph_; tools::graph<ir::value*> graph_;
std::map<ir::value*, size_t> groups_; std::map<ir::value*, size_t> groups_;
std::map<size_t, std::vector<ir::value*>> values_; std::map<size_t, std::vector<ir::value*>> values_;
std::map<size_t, layout_t> layouts_; std::map<size_t, layout_t*> layouts_;
}; };
} }

4
include/triton/codegen/analysis/liveness.h
View File

@@ -45,7 +45,9 @@ struct double_buffer_info_t {
struct buffer_t { struct buffer_t {
size_t id; size_t id;
size_t size; size_t size;
bool operator<(buffer_t other) const { return id < other.id; } bool operator<(buffer_t other) const {
return id < other.id;
}
}; };
class liveness { class liveness {

2
include/triton/codegen/selection.h
View File

@@ -62,7 +62,7 @@ class target;
typedef std::vector<Value*> indices_t; typedef std::vector<Value*> indices_t;
struct distributed_axis { struct distributed_axis {
size_t contiguous; int contiguous;
std::vector<Value*> values; std::vector<Value*> values;
Value* thread_id; Value* thread_id;
}; };

1
include/triton/ir/type.h
View File

@@ -64,6 +64,7 @@ public:
type *get_scalar_ty() const; type *get_scalar_ty() const;
const tile_shapes_t& get_tile_shapes() const; const tile_shapes_t& get_tile_shapes() const;
const size_t get_tile_rank() const; const size_t get_tile_rank() const;
const size_t get_tile_ranks1() const;
unsigned get_tile_num_elements() const; unsigned get_tile_num_elements() const;
type *get_tile_element_ty() const; type *get_tile_element_ty() const;
unsigned get_pointer_address_space() const; unsigned get_pointer_address_space() const;

171
lib/codegen/analysis/layout.cc
View File

@@ -58,6 +58,7 @@ void layout::make_graph(ir::instruction *i) {
} }
} }
// hmma // hmma
bool is_hmma_c(ir::value *v){ bool is_hmma_c(ir::value *v){
bool result = false; bool result = false;
@@ -72,11 +73,11 @@ bool is_hmma_c(ir::value *v){
return result; return result;
} }
const layout_t &layout::get(ir::value *v) const { const layout_t* layout::get(ir::value *v) const {
return layouts_.at(groups_.at(v)); return layouts_.at(groups_.at(v));
} }
std::map<size_t, layout_t>& layout::get_all() { std::map<size_t, layout_t*>& layout::get_all() {
return layouts_; return layouts_;
} }
@@ -102,19 +103,45 @@ inline bool is_trans(ir::value *v) {
return false; return false;
} }
layout_t::layout_t(layout_type_t _type,
const std::vector<int> &_axes,
const std::vector<unsigned> &_shapes,
const std::vector<ir::value *> &values,
analysis::align* align): type(_type), axes(_axes), shapes(_shapes) {
// io pointer
std::set<ir::value*> ptr;
for(ir::value* v: values)
extract_io_use(v, ptr);
size_t rank = axes.size();
std::vector<int> order(rank);
std::iota(order.begin(), order.end(), 0);
for(ir::value *v: ptr){
auto max_contiguous = align->contiguous(v);
std::sort(order.begin(), order.end(), [&](unsigned a, unsigned b) {
return max_contiguous[a] > max_contiguous[b];
});
}
this->order = order;
}
inline unsigned clamp(unsigned x, unsigned lo, unsigned hi) { inline unsigned clamp(unsigned x, unsigned lo, unsigned hi) {
return std::min(std::max(x, lo), hi); return std::min(std::max(x, lo), hi);
} }
void layout::init_hmma_tile(layout_t& layout) { layout_hmma_884_t::layout_hmma_884_t(size_t num_warps,
auto ord = layout.order; const std::vector<int>& _axes,
auto shapes = layout.shapes; const std::vector<unsigned>& _shapes,
unsigned shape_0 = shapes[ord[0]]; const std::vector<ir::value *> &values,
unsigned shape_1 = shapes[ord[1]]; analysis::align* align): layout_t(HMMA_884, _axes, _shapes, values, align) {
unsigned shape_0 = shapes[order[0]];
unsigned shape_1 = shapes[order[1]];
/* fragments per warp */ /* fragments per warp */
// try to make things as square as possible to maximize data re-use // try to make things as square as possible to maximize data re-use
std::vector<unsigned> fpw = {1, 1, 1}; fpw = {1, 1, 1};
std::vector<unsigned> fpw_nm1; std::vector<int> fpw_nm1;
unsigned num_fragments = std::min<unsigned>((shape_0/8)*(shape_1/8), 4); unsigned num_fragments = std::min<unsigned>((shape_0/8)*(shape_1/8), 4);
do { do {
fpw_nm1 = fpw; fpw_nm1 = fpw;
@@ -123,144 +150,108 @@ void layout::init_hmma_tile(layout_t& layout) {
if(fpw[0]*fpw[1] < num_fragments) if(fpw[0]*fpw[1] < num_fragments)
fpw[1] = clamp(fpw[1]*2, 1, shape_1 / 8); fpw[1] = clamp(fpw[1]*2, 1, shape_1 / 8);
}while(fpw_nm1 != fpw); }while(fpw_nm1 != fpw);
// store parameters
for(unsigned d = 0; d < shapes.size(); d++)
layout.fpw[d] = fpw[d];
/* warps per tile */ /* warps per tile */
// try to make things as square as possible to maximize data re-use // try to make things as square as possible to maximize data re-use
std::vector<unsigned> wpt = {1, 1, 1}; wpt = {1, 1, 1};
std::vector<unsigned> wpt_nm1; std::vector<int> wpt_nm1;
do{ do{
wpt_nm1 = wpt; wpt_nm1 = wpt;
if(wpt[0] * wpt[1] * wpt[2] < num_warps_) if(wpt[0] * wpt[1] * wpt[2] < num_warps)
wpt[0] = clamp(wpt[0]*2, 1, shape_0 / (fpw[0]*8)); wpt[0] = clamp(wpt[0]*2, 1, shape_0 / (fpw[0]*8));
if(wpt[0] * wpt[1] * wpt[2] < num_warps_) if(wpt[0] * wpt[1] * wpt[2] < num_warps)
wpt[1] = clamp(wpt[1]*2, 1, shape_1 / (fpw[1]*8)); wpt[1] = clamp(wpt[1]*2, 1, shape_1 / (fpw[1]*8));
}while(wpt_nm1 != wpt); }while(wpt_nm1 != wpt);
// store parameters
for(unsigned d = 0; d < shapes.size(); d++)
layout.wpt[d] = wpt[d];
/* sanity check */ /* sanity check */
unsigned effective_num_warps = 1; unsigned effective_num_warps = 1;
for(size_t d = 0; d < shapes.size(); d++) for(size_t d = 0; d < shapes.size(); d++)
effective_num_warps *= layout.wpt[d]; effective_num_warps *= wpt[d];
if(num_warps_ != effective_num_warps) if(num_warps != effective_num_warps)
throw std::runtime_error("cannot create a kernel with this amount of warps"); throw std::runtime_error("cannot create a kernel with this amount of warps");
} }
void layout::init_scanline_tile(layout_t& layout) { layout_scanline_t::layout_scanline_t(size_t num_warps,
auto ord = layout.order; const std::vector<int>& _axes,
auto shapes = layout.shapes; const std::vector<unsigned>& _shapes,
const std::vector<ir::value *> &values,
analysis::align* align): layout_t(SCANLINE, _axes, _shapes, values, align){
unsigned size = std::accumulate(shapes.begin(), shapes.end(), 1, std::multiplies<int>()); unsigned size = std::accumulate(shapes.begin(), shapes.end(), 1, std::multiplies<int>());
unsigned ld = ord[0]; unsigned num_threads = num_warps * 32;
unsigned num_threads = num_warps_*32; nts.resize(shapes.size());
unsigned current = num_threads; mts.resize(shapes.size());
layout.nts[ld] = clamp(size / num_threads, 1, 4); unsigned i = order[0];
layout.mts[ld] = clamp(current, 1, shapes[ld] / layout.nts[ld]); nts[i] = clamp(size / num_threads, 1, 4);
current = current / layout.mts[ld]; mts[i] = clamp(num_threads, 1, shapes[i] / nts[i]);
num_threads = num_threads / mts[i];
for(size_t d = 1; d < shapes.size(); d++){ for(size_t d = 1; d < shapes.size(); d++){
ld = ord[d]; i = order[d];
layout.nts[ld] = 1; nts[i] = 1;
layout.mts[ld] = clamp(current, 1, shapes[ld]); mts[i] = clamp(num_threads, 1, shapes[i]);
current = current / layout.mts[ld]; num_threads = num_threads / mts[i];
} }
/* sanity check */ /* sanity check */
unsigned effective_num_threads = 1; unsigned effective_num_threads = 1;
for(size_t d = 0; d < shapes.size(); d++) for(size_t d = 0; d < shapes.size(); d++)
effective_num_threads *= layout.mts[d]; effective_num_threads *= mts[d];
if(num_threads != effective_num_threads) if(num_warps * 32 != effective_num_threads)
throw std::runtime_error("cannot create a kernel with this amount of warps"); throw std::runtime_error("cannot create a kernel with this amount of warps");
} }
void layout::run(ir::module &mod) { void layout::run(ir::module &mod) {
// make graph // make graph
graph_.clear(); graph_.clear();
ir::for_each_instruction(mod, [this](ir::instruction* i) { ir::for_each_instruction(mod, [this](ir::instruction* i) {
make_graph(i); make_graph(i);
}); });
// connected components // connected components
graph_.connected_components(&values_, &groups_); graph_.connected_components(&values_, &groups_);
// create layouts // create layouts
for(const auto& x: values_) { for(const auto& x: values_) {
bool hmma_c = std::any_of(x.second.begin(), x.second.end(), &is_hmma_c); bool hmma_c = std::any_of(x.second.begin(), x.second.end(), &is_hmma_c);
// type auto cmp = [](ir::value* x, ir::value *y) {
layouts_[x.first].type = hmma_c ? HMMA_884 : SCANLINE; return x->get_type()->get_tile_ranks1() <
} y->get_type()->get_tile_ranks1();
};
/* ---- TO CLEAN ---- */
size_t num_groups = num_layouts();
// helpers
auto rank = [this](ir::value* v) {
int ret = 0;
for(int s: v->get_type()->get_tile_shapes())
ret += s > 1;
return ret;
};
// find out axes for each layout
for(const auto& x: values_) {
auto cmp = [&rank](ir::value* x, ir::value *y) { return rank(x) < rank(y); };
ir::value *largest = *std::max_element(x.second.begin(), x.second.end(), cmp); ir::value *largest = *std::max_element(x.second.begin(), x.second.end(), cmp);
layouts_[x.first].axes = axes_->get(largest); const auto& axes = axes_->get(largest);
layouts_[x.first].shapes = largest->get_type()->get_tile_shapes(); const auto& shapes = largest->get_type()->get_tile_shapes();
// type
if(hmma_c)
layouts_[x.first] = new layout_hmma_884_t(num_warps_, axes, shapes, x.second, align_);
else
layouts_[x.first] = new layout_scanline_t(num_warps_, axes, shapes, x.second, align_);
} }
// find out the layout ordering of a group
for(const auto& x: values_) {
std::set<ir::value*> ptr;
for(ir::value* v: x.second)
extract_io_use(v, ptr);
size_t rank = layouts_[x.first].axes.size();
std::vector<int> order(rank);
std::iota(order.begin(), order.end(), 0);
for(ir::value *v: ptr){
auto max_contiguous = align_->contiguous(v);
std::sort(order.begin(), order.end(), [&](unsigned a, unsigned b) {
return max_contiguous[a] > max_contiguous[b]; }
);
}
layouts_[x.first].order = order;
}
// matrix multiplication optimizations // matrix multiplication optimizations
for(size_t i = 0; i < num_groups; i++){ for(const auto& x: values_) {
std::vector<ir::dot_inst*> dots; std::vector<ir::dot_inst*> dots;
for(ir::value* v: values_of(i)) for(ir::value* v: x.second)
if(auto *x = dynamic_cast<ir::dot_inst*>(v)) if(auto *x = dynamic_cast<ir::dot_inst*>(v))
dots.push_back(x); dots.push_back(x);
for(ir::dot_inst* dot: dots){ for(ir::dot_inst* dot: dots){
ir::value* a = dot->get_operand(0); ir::value* a = dot->get_operand(0);
ir::value* b = dot->get_operand(1); ir::value* b = dot->get_operand(1);
if(get(dot).type == HMMA_884){ if(get(dot)->type == HMMA_884){
auto a_val = values_of(layout_of(a)); auto a_val = values_of(layout_of(a));
auto b_val = values_of(layout_of(b)); auto b_val = values_of(layout_of(b));
for(ir::value *v: a_val) for(ir::value *v: a_val)
if(auto *cts = dynamic_cast<ir::copy_to_shared_inst*>(v)) if(auto *cts = dynamic_cast<ir::copy_to_shared_inst*>(v))
layouts_[layout_of(a)].order = layouts_[layout_of(cts->get_operand(0))].order; layouts_[layout_of(a)]->order = layouts_[layout_of(cts->get_operand(0))]->order;
for(ir::value *v: b_val) for(ir::value *v: b_val)
if(auto *cts = dynamic_cast<ir::copy_to_shared_inst*>(v)) if(auto *cts = dynamic_cast<ir::copy_to_shared_inst*>(v))
layouts_[layout_of(b)].order = layouts_[layout_of(cts->get_operand(0))].order; layouts_[layout_of(b)]->order = layouts_[layout_of(cts->get_operand(0))]->order;
} }
else{ else{
std::vector<int> col = {0, 1}; std::vector<int> col = {0, 1};
std::vector<int> row = {1, 0}; std::vector<int> row = {1, 0};
layouts_[layout_of(a)].order = is_trans(a) ? row : col; layouts_[layout_of(a)]->order = is_trans(a) ? row : col;
layouts_[layout_of(b)].order = is_trans(b) ? col : row; layouts_[layout_of(b)]->order = is_trans(b) ? col : row;
} }
} }
} }
// tiling parameters
for(auto& x: layouts_){
/* HMMA parameters*/
if(x.second.type == HMMA_884)
init_hmma_tile(x.second);
else
init_scanline_tile(x.second);
}
} }
} }

26
lib/codegen/analysis/liveness.cc
View File

@@ -89,8 +89,8 @@ bool liveness::do_pad(ir::value *x) {
ir::value *b = dot->get_operand(1); ir::value *b = dot->get_operand(1);
size_t a_previous = pad_[a]; size_t a_previous = pad_[a];
size_t b_previous = pad_[b]; size_t b_previous = pad_[b];
auto a_order = layouts_->get(a).order; auto a_order = layouts_->get(a)->order;
auto b_order = layouts_->get(b).order; auto b_order = layouts_->get(b)->order;
bool a_row = is_trans(a) ^ (a_order[0] == 1); bool a_row = is_trans(a) ^ (a_order[0] == 1);
bool b_row = is_trans(b) ^ (b_order[0] == 1); bool b_row = is_trans(b) ^ (b_order[0] == 1);
auto a_shapes = a->get_type()->get_tile_shapes(); auto a_shapes = a->get_type()->get_tile_shapes();
@@ -108,9 +108,9 @@ bool liveness::do_pad(ir::value *x) {
} }
// padding for copy to shared // padding for copy to shared
if(auto* cts = dynamic_cast<ir::copy_to_shared_inst*>(x)) { if(auto* cts = dynamic_cast<ir::copy_to_shared_inst*>(x)) {
auto cts_order = layouts_->get(cts).order; auto cts_order = layouts_->get(cts)->order;
ir::value *arg = cts->get_operand(0); ir::value *arg = cts->get_operand(0);
auto arg_order = layouts_->get(arg).order; auto arg_order = layouts_->get(arg)->order;
size_t previous = pad_[cts]; size_t previous = pad_[cts];
if(cts_order != arg_order) if(cts_order != arg_order)
pad_[cts] = std::max<int>(pad_[cts], 4); pad_[cts] = std::max<int>(pad_[cts], 4);
@@ -134,26 +134,10 @@ bool liveness::do_pad(ir::value *x) {
} }
unsigned liveness::num_bytes(ir::value *x) { unsigned liveness::num_bytes(ir::value *x) {
if(auto *red = dynamic_cast<ir::reduce_inst*>(x)){
unsigned num_bytes = x->get_type()->get_scalar_ty()->get_primitive_size_in_bits() / 8;
size_t axis = red->get_axis();
ir::value *op = red->get_operand(0);
auto shapes = op->get_type()->get_tile_shapes();
shapes.erase(shapes.begin() + axis);
size_t num_elements = 1;
for(auto x: shapes)
num_elements *= x;
size_t depth;
if(layouts_->get(x).type == HMMA_884)
depth = layouts_->get(op).wpt.at(axis);
else
depth = layouts_->get(op).mts.at(axis);
return num_elements * num_bytes * depth;
}
unsigned num_bytes = x->get_type()->get_primitive_size_in_bits() / 8; unsigned num_bytes = x->get_type()->get_primitive_size_in_bits() / 8;
unsigned pad = pad_.at(x); unsigned pad = pad_.at(x);
if(pad > 0){ if(pad > 0){
unsigned ld = x->get_type()->get_tile_shapes()[layouts_->get(x).order[0]]; unsigned ld = x->get_type()->get_tile_shapes()[layouts_->get(x)->order[0]];
num_bytes += pad * num_bytes / ld; num_bytes += pad * num_bytes / ld;
} }
if(has_double(x)) if(has_double(x))

43
lib/codegen/selection.cc
View File

@@ -559,7 +559,7 @@ Value* selection::llvm_value(ir::value *v, IRBuilder<> &builder) {
* ------------------- */ * ------------------- */
// Grid construction // Grid construction
std::vector<Value*> delinearize(Value *trailing, const std::vector<int>& order, std::vector<unsigned> &shapes, IRBuilder<> &builder){ std::vector<Value*> delinearize(Value *trailing, const std::vector<int>& order, std::vector<int> &shapes, IRBuilder<> &builder){
size_t dim = shapes.size(); size_t dim = shapes.size();
std::vector<Value*> result(dim); std::vector<Value*> result(dim);
for(unsigned k = 0; k < dim - 1; k++){ for(unsigned k = 0; k < dim - 1; k++){
@@ -580,12 +580,8 @@ void selection::init_strided_scan_axes(const analysis::layout_t& layout, IRBuild
auto order = layout.order; auto order = layout.order;
const auto& shapes = layout.shapes; const auto& shapes = layout.shapes;
size_t dim = shapes.size(); size_t dim = shapes.size();
std::vector<unsigned> nts(dim); std::vector<int> nts = layout.nts;
std::vector<unsigned> mts(dim); std::vector<int> mts = layout.mts;
for(unsigned i = 0; i < shapes.size(); i++){
nts[i] = layout.nts.at(i);
mts[i] = layout.mts.at(i);
}
Value* full_thread_id = builder.CreateAdd(builder.CreateMul(u_warp_id, builder.getInt32(32)), u_thread_id); Value* full_thread_id = builder.CreateAdd(builder.CreateMul(u_warp_id, builder.getInt32(32)), u_thread_id);
std::vector<Value*> thread_id = delinearize(full_thread_id, order, mts, builder); std::vector<Value*> thread_id = delinearize(full_thread_id, order, mts, builder);
// Create axes // Create axes
@@ -608,6 +604,7 @@ void selection::init_hmma_axes(const analysis::layout_t& layout, IRBuilder<> &bu
const auto& shapes = layout.shapes; const auto& shapes = layout.shapes;
if(shapes.size() > 3) if(shapes.size() > 3)
throw std::runtime_error("unsupported"); throw std::runtime_error("unsupported");
bool is_batched = shapes.size() >= 3; bool is_batched = shapes.size() >= 3;
Value *_1 = builder.getInt32(1); Value *_1 = builder.getInt32(1);
@@ -725,7 +722,7 @@ void selection::init_axes(const analysis::layout_t& layout, IRBuilder<> &builder
void selection::create_shared_tile(ir::value *v, IRBuilder<> &builder, Value *sh_mem_ptr) { void selection::create_shared_tile(ir::value *v, IRBuilder<> &builder, Value *sh_mem_ptr) {
if(tmap_.find(v) != tmap_.end()) if(tmap_.find(v) != tmap_.end())
return; return;
auto order = layouts_->get(v).order; auto order = layouts_->get(v)->order;
auto shapes = v->get_type()->get_tile_shapes(); auto shapes = v->get_type()->get_tile_shapes();
unsigned pad = liveness_->get_pad(v); unsigned pad = liveness_->get_pad(v);
if(pad > 0) if(pad > 0)
@@ -775,7 +772,7 @@ void selection::create_distributed_tile(ir::value *v, IRBuilder<> &builder) {
axes[d].values = {builder.getInt32(0)}; axes[d].values = {builder.getInt32(0)};
} }
} }
distributed_tile *T = new distributed_tile(ty, shapes, layouts_->get(v).order, axes, builder, false); distributed_tile *T = new distributed_tile(ty, shapes, layouts_->get(v)->order, axes, builder, false);
bool is_inserted = tmap_.insert({v, T}).second; bool is_inserted = tmap_.insert({v, T}).second;
// constant range // constant range
if(is_inserted && dynamic_cast<ir::make_range*>(v)){ if(is_inserted && dynamic_cast<ir::make_range*>(v)){
@@ -819,7 +816,7 @@ void selection::init_layouts(ir::function *fn, IRBuilder<> &builder, Value *sh_m
Value *u_warp_id = builder.CreateUDiv(u_thread_id, warp_size); Value *u_warp_id = builder.CreateUDiv(u_thread_id, warp_size);
// create grid // create grid
for(auto x: layouts_->get_all()) for(auto x: layouts_->get_all())
init_axes(x.second, builder, u_thread_warp_id, u_warp_id); init_axes(*x.second, builder, u_thread_warp_id, u_warp_id);
// create tile // create tile
std::set<ir::value*> seen; std::set<ir::value*> seen;
for(ir::basic_block *block: fn->blocks()) for(ir::basic_block *block: fn->blocks())
@@ -932,7 +929,7 @@ void selection::lower_reduce(ir::reduce_inst *x, LLVMContext &ctx, Function *fn,
tgt_->add_barrier(module, builder); tgt_->add_barrier(module, builder);
builder.CreateStore(result, write_ptr); builder.CreateStore(result, write_ptr);
// build result // build result
unsigned depth = layouts_->get(op).wpt.at(axis); unsigned depth = layouts_->get(op)->wpt.at(axis);
for(unsigned i = depth/2; i > 0; i >>= 1){ for(unsigned i = depth/2; i > 0; i >>= 1){
// current indices // current indices
indices_t current(write_idx.size(), builder.getInt32(0)); indices_t current(write_idx.size(), builder.getInt32(0));
@@ -1013,15 +1010,15 @@ void selection::lower_broadcast(ir::broadcast_inst *x, LLVMContext &ctx, Functio
void selection::lower_copy_to_shared(ir::copy_to_shared_inst *x, LLVMContext &ctx, Function *fn, IRBuilder<> &builder) { void selection::lower_copy_to_shared(ir::copy_to_shared_inst *x, LLVMContext &ctx, Function *fn, IRBuilder<> &builder) {
unsigned vector_size = 1; unsigned vector_size = 1;
auto x_order = layouts_->get(x).order; auto x_order = layouts_->get(x)->order;
ir::value *arg = x->get_operand(0); ir::value *arg = x->get_operand(0);
auto arg_order = layouts_->get(arg).order; auto arg_order = layouts_->get(arg)->order;
// tiles // tiles
shared_tile* result = (shared_tile*)tmap_.at(x); shared_tile* result = (shared_tile*)tmap_.at(x);
distributed_tile* in = (distributed_tile*)tmap_.at(arg); distributed_tile* in = (distributed_tile*)tmap_.at(arg);
if(x_order == arg_order){ if(x_order == arg_order){
size_t ld = arg_order[0]; size_t ld = arg_order[0];
vector_size = std::min(layouts_->get(x).nts.at(ld), layouts_->get(arg).nts.at(ld)); vector_size = std::min(layouts_->get(x)->nts.at(ld), layouts_->get(arg)->nts.at(ld));
} }
std::map<unsigned, Value*> packets; std::map<unsigned, Value*> packets;
@@ -1090,8 +1087,8 @@ void selection::lower_hmma_dot(ir::dot_inst *dot, LLVMContext &ctx, Function *fn
Value* u_thread_id = tgt_->get_local_id(builder.GetInsertBlock()->getModule(), builder, 0); Value* u_thread_id = tgt_->get_local_id(builder.GetInsertBlock()->getModule(), builder, 0);
auto ord_a = layouts_->get(dot->get_operand(0)).order; auto ord_a = layouts_->get(dot->get_operand(0))->order;
auto ord_b = layouts_->get(dot->get_operand(1)).order; auto ord_b = layouts_->get(dot->get_operand(1))->order;
bool is_a_trans = is_trans(dot->get_operand(0)); bool is_a_trans = is_trans(dot->get_operand(0));
bool is_b_trans = is_trans(dot->get_operand(1)); bool is_b_trans = is_trans(dot->get_operand(1));
@@ -1117,12 +1114,12 @@ void selection::lower_hmma_dot(ir::dot_inst *dot, LLVMContext &ctx, Function *fn
"{$10, $11}, " "{$10, $11}, "
"{$0, $1, $2, $3, $4, $5, $6, $7};", "=f,=f,=f,=f,=f,=f,=f,=f,r,r,r,r,0,1,2,3,4,5,6,7", false); "{$0, $1, $2, $3, $4, $5, $6, $7};", "=f,=f,=f,=f,=f,=f,=f,=f,r,r,r,r,0,1,2,3,4,5,6,7", false);
unsigned fpw_0 = layouts_->get(dot).fpw.at(0); unsigned fpw_0 = layouts_->get(dot)->fpw.at(0);
unsigned fpw_1 = layouts_->get(dot).fpw.at(1); unsigned fpw_1 = layouts_->get(dot)->fpw.at(1);
unsigned wts_0 = fpw_0 * 8; unsigned wts_0 = fpw_0 * 8;
unsigned wts_1 = fpw_1 * 8; unsigned wts_1 = fpw_1 * 8;
unsigned wpt_0 = layouts_->get(dot).wpt.at(0); unsigned wpt_0 = layouts_->get(dot)->wpt.at(0);
unsigned wpt_1 = layouts_->get(dot).wpt.at(1); unsigned wpt_1 = layouts_->get(dot)->wpt.at(1);
unsigned stride_rep_i = wpt_0 * wts_0; unsigned stride_rep_i = wpt_0 * wts_0;
unsigned stride_rep_j = wpt_1 * wts_1; unsigned stride_rep_j = wpt_1 * wts_1;
unsigned num_rep_i = shapes[0] / stride_rep_i; unsigned num_rep_i = shapes[0] / stride_rep_i;
@@ -1253,7 +1250,7 @@ void selection::lower_dot(ir::dot_inst *dot, LLVMContext &ctx, Function *fn, IRB
if(NK != 1) { if(NK != 1) {
shared_tile *TA = (shared_tile*)tmap_.at(A); shared_tile *TA = (shared_tile*)tmap_.at(A);
shared_tile *TB = (shared_tile*)tmap_.at(B); shared_tile *TB = (shared_tile*)tmap_.at(B);
if(layouts_->get(dot).type == analysis::HMMA_884) if(layouts_->get(dot)->type == analysis::HMMA_884)
lower_hmma_dot(dot, ctx, fn, builder, TC, TA, TB, TD, NK); lower_hmma_dot(dot, ctx, fn, builder, TC, TA, TB, TD, NK);
else else
lower_scanline_dot(dot, ctx, fn, builder, TC, TA, TB, TD, NK, c_ty, f_mul_add); lower_scanline_dot(dot, ctx, fn, builder, TC, TA, TB, TD, NK, c_ty, f_mul_add);
@@ -1269,7 +1266,7 @@ void selection::lower_masked_load(ir::masked_load_inst *x, LLVMContext &ctx, Fun
// find vector size // find vector size
distributed_tile* result = (distributed_tile*)tmap_.at(x); distributed_tile* result = (distributed_tile*)tmap_.at(x);
ir::value *ptr = x->get_pointer_operand(); ir::value *ptr = x->get_pointer_operand();
size_t ld = layouts_->get(ptr).order[0]; size_t ld = layouts_->get(ptr)->order[0];
unsigned alignment = alignment_->get(ptr, ld); unsigned alignment = alignment_->get(ptr, ld);
unsigned vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment); unsigned vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment);
distributed_tile *pointers = (distributed_tile*)tmap_.at(ptr); distributed_tile *pointers = (distributed_tile*)tmap_.at(ptr);
@@ -1341,7 +1338,7 @@ void selection::lower_load(ir::load_inst *x, LLVMContext &ctx, Function *fn, IRB
distributed_tile* result = (distributed_tile*)tmap_.at(x); distributed_tile* result = (distributed_tile*)tmap_.at(x);
// find vector size // find vector size
ir::value *ptr = x->get_pointer_operand(); ir::value *ptr = x->get_pointer_operand();
size_t ld = layouts_->get(ptr).order[0]; size_t ld = layouts_->get(ptr)->order[0];
unsigned alignment = alignment_->get(ptr, ld); unsigned alignment = alignment_->get(ptr, ld);
unsigned vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment); unsigned vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment);
distributed_tile *pointers = (distributed_tile*)tmap_.at(ptr); distributed_tile *pointers = (distributed_tile*)tmap_.at(ptr);

1
lib/codegen/transform/coalesce.cc
View File

@@ -53,7 +53,6 @@ ir::value* coalesce::rematerialize(ir::value *x, ir::builder &builder,
builder.set_insert_point(pos); builder.set_insert_point(pos);
if(dynamic_cast<ir::load_inst*>(x)){ if(dynamic_cast<ir::load_inst*>(x)){
ir::value *ret = builder.insert(ir::copy_to_shared_inst::create(x)); ir::value *ret = builder.insert(ir::copy_to_shared_inst::create(x));
// x->replace_all_uses_with(ret);
return ret; return ret;
} }
// default -- recursive clone // default -- recursive clone

8
lib/ir/type.cc
View File

@@ -77,6 +77,14 @@ const size_t type::get_tile_rank() const {
return get_tile_shapes().size(); return get_tile_shapes().size();
} }
const size_t type::get_tile_ranks1() const {
int ret = 0;
for(int s: get_tile_shapes())
ret += s > 1;
return ret;
}
unsigned type::get_tile_num_elements() const { unsigned type::get_tile_num_elements() const {
const tile_shapes_t& shapes = get_tile_shapes(); const tile_shapes_t& shapes = get_tile_shapes();
unsigned result = 1; unsigned result = 1;