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

[code generation] now vectorizing shared memory stores

Browse Source
This commit is contained in:
Philippe Tillet
2019-02-10 21:59:41 -05:00
parent 8ab5ca3de3
commit b2e487491f
8 changed files with 64 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
include/codegen/selection.h
View File

@@ -68,42 +68,25 @@ class distributed_tile: public tile{
private:
void init_indices();
public:
distributed_tile(llvm::Type *ty, const shapes_t& shapes, const axes_t &axes);
virtual void for_each(std::function<void(indices_t)> fn) = 0;
protected:
axes_t axes_;
indices_map_t indices_;
values_t values_;
};
class serialized_distributed_tile: public distributed_tile {
public:
using distributed_tile::distributed_tile;
public:
void set_value(indices_t, llvm::Value *);
llvm::Value* get_value(indices_t idx);
void for_each(std::function<void(indices_t)> fn);
};
class vectorized_distributed_tile: public distributed_tile {
private:
llvm::Type *make_vector_ty(llvm::Type *ty, size_t vector_size);
public:
vectorized_distributed_tile(llvm::Type *ty, const shapes_t& shapes, const axes_t &axes, llvm::IRBuilder<> &builder);
void set_value(indices_t, llvm::Value *);
distributed_tile(llvm::Type *ty, const shapes_t& shapes, const axes_t &axes, llvm::IRBuilder<> &builder, bool vectorize);
void set_value(indices_t idx, llvm::Value *v);
llvm::Value* get_value(indices_t idx);
unsigned get_linear_index(indices_t idx);
void for_each(std::function<void(indices_t)> fn);
const distributed_axis &axis(unsigned dim) { return axes_.at(dim); }
private:
llvm::IRBuilder<> &builder_;
axes_t axes_;
indices_map_t indices_;
values_t values_;
size_t vector_size_;
llvm::IRBuilder<> &builder_;
};
class selection{
typedef std::map<ir::value *, llvm::Value *> vmap_t;
typedef std::map<ir::value *, tile *> tmap_t;

1
include/codegen/tune.h
View File

@@ -32,6 +32,7 @@ public:
std::vector<unsigned *> get_params(ir::module& mod);
std::map<std::string, unsigned *> get_params(ir::instruction* i);
unsigned *get_param(ir::value *value, const std::string &key) { return params_[value][key]; }
void copy(ir::value *dst, ir::value *src) { params_[dst] = params_[src]; }
bool check_constraints(ir::module &fn, std::map<ir::value *, std::vector<std::string>> &errors);
void run(ir::module &mod);