[codegen] more cleaning

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

View File

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

View File

@@ -45,7 +45,9 @@ struct double_buffer_info_t {
struct buffer_t {
size_t id;
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 {

View File

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

View File

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

View File

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

View File

@@ -89,8 +89,8 @@ bool liveness::do_pad(ir::value *x) {
ir::value *b = dot->get_operand(1);
size_t a_previous = pad_[a];
size_t b_previous = pad_[b];
auto a_order = layouts_->get(a).order;
auto b_order = layouts_->get(b).order;
auto a_order = layouts_->get(a)->order;
auto b_order = layouts_->get(b)->order;
bool a_row = is_trans(a) ^ (a_order[0] == 1);
bool b_row = is_trans(b) ^ (b_order[0] == 1);
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
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);
auto arg_order = layouts_->get(arg).order;
auto arg_order = layouts_->get(arg)->order;
size_t previous = pad_[cts];
if(cts_order != arg_order)
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) {
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 pad = pad_.at(x);
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;
}
if(has_double(x))

View File

@@ -559,7 +559,7 @@ Value* selection::llvm_value(ir::value *v, IRBuilder<> &builder) {
* ------------------- */
// 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();
std::vector<Value*> result(dim);
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;
const auto& shapes = layout.shapes;
size_t dim = shapes.size();
std::vector<unsigned> nts(dim);
std::vector<unsigned> mts(dim);
for(unsigned i = 0; i < shapes.size(); i++){
nts[i] = layout.nts.at(i);
mts[i] = layout.mts.at(i);
}
std::vector<int> nts = layout.nts;
std::vector<int> mts = layout.mts;
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);
// Create axes
@@ -608,6 +604,7 @@ void selection::init_hmma_axes(const analysis::layout_t& layout, IRBuilder<> &bu
const auto& shapes = layout.shapes;
if(shapes.size() > 3)
throw std::runtime_error("unsupported");
bool is_batched = shapes.size() >= 3;
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) {
if(tmap_.find(v) != tmap_.end())
return;
auto order = layouts_->get(v).order;
auto order = layouts_->get(v)->order;
auto shapes = v->get_type()->get_tile_shapes();
unsigned pad = liveness_->get_pad(v);
if(pad > 0)
@@ -775,7 +772,7 @@ void selection::create_distributed_tile(ir::value *v, IRBuilder<> &builder) {
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;
// constant range
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);
// create grid
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
std::set<ir::value*> seen;
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);
builder.CreateStore(result, write_ptr);
// 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){
// current indices
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) {
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);
auto arg_order = layouts_->get(arg).order;
auto arg_order = layouts_->get(arg)->order;
// tiles
shared_tile* result = (shared_tile*)tmap_.at(x);
distributed_tile* in = (distributed_tile*)tmap_.at(arg);
if(x_order == arg_order){
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;
@@ -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);
auto ord_a = layouts_->get(dot->get_operand(0)).order;
auto ord_b = layouts_->get(dot->get_operand(1)).order;
auto ord_a = layouts_->get(dot->get_operand(0))->order;
auto ord_b = layouts_->get(dot->get_operand(1))->order;
bool is_a_trans = is_trans(dot->get_operand(0));
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}, "
"{$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_1 = layouts_->get(dot).fpw.at(1);
unsigned fpw_0 = layouts_->get(dot)->fpw.at(0);
unsigned fpw_1 = layouts_->get(dot)->fpw.at(1);
unsigned wts_0 = fpw_0 * 8;
unsigned wts_1 = fpw_1 * 8;
unsigned wpt_0 = layouts_->get(dot).wpt.at(0);
unsigned wpt_1 = layouts_->get(dot).wpt.at(1);
unsigned wpt_0 = layouts_->get(dot)->wpt.at(0);
unsigned wpt_1 = layouts_->get(dot)->wpt.at(1);
unsigned stride_rep_i = wpt_0 * wts_0;
unsigned stride_rep_j = wpt_1 * wts_1;
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) {
shared_tile *TA = (shared_tile*)tmap_.at(A);
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);
else
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
distributed_tile* result = (distributed_tile*)tmap_.at(x);
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 vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment);
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);
// find vector size
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 vector_size = std::min<unsigned>(result->axis(ld).contiguous, alignment);
distributed_tile *pointers = (distributed_tile*)tmap_.at(ptr);

View File

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

View File

@@ -77,6 +77,14 @@ const size_t type::get_tile_rank() const {
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 {
const tile_shapes_t& shapes = get_tile_shapes();
unsigned result = 1;