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

matmul test passes

Browse Source
This commit is contained in:
Philippe Tillet
2019-08-23 17:13:30 -07:00
parent 64a6910644
commit 8798d240dc
5 changed files with 122 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
examples/cpp/dot.cpp
View File

@@ -82,13 +82,8 @@ R"(
#define true 1
#define false 0
#define __bool_true_false_are_defined 1
extern int get_program_id(int);
static const int TM __attribute__((one_of(128)));
static const int TN __attribute__((one_of(128)));
static const int TK __attribute__((one_of(32)));
void matmul(restrict )" + a_ty + R"( * A __attribute__((readonly, aligned(16))),
restrict )" + b_ty + R"( * B __attribute__((readonly, aligned(16))),
restrict )" + c_ty + R"( * C __attribute__((aligned(16))),
@@ -162,10 +157,15 @@ perf_t do_bench(drv::stream* stream, bool AT, bool BT, int32_t M, int32_t N, int
stream->write(dc, true, 0, hc);
stream->synchronize();
// run
rt::function function(src(AT, BT, ty, ty, ty, 8, 8));
rt::function::options_space_t opt;
opt.defines.push_back({"TM", {"128"}});
opt.defines.push_back({"TN", {"128"}});
opt.defines.push_back({"TK", {"32"}});
opt.num_warps = {1, 2, 4, 8};
rt::function function(src(AT, BT, ty, ty, ty, 8, 8), opt);
auto ceil = [](size_t x, size_t y) { return (x + y - 1) / y; };
auto grid = [&](const rt::params_t& x) { return rt::grid_t{ceil(M, x.at("TM")), ceil(N, x.at("TN")), 1}; };
auto grid = [&](const rt::function::options_t& x) { return rt::grid_t{ceil(M, x.D<int>("TM")), ceil(N, x.D<int>("TN")), 1}; };
auto tflops = [&](double nanosec) { return 2.*M*N*K / nanosec * 1e-3; };
perf_t res;

45
include/triton/runtime/function.h
View File

@@ -60,54 +60,69 @@ namespace runtime{
typedef std::array<size_t, 3> grid_t;
typedef std::map<std::string, size_t> params_t;
struct options {
size_t num_warps;
params_t params;
};
template<typename T> T convert(const std::string& name);
template<> long convert<long>(const std::string& name) { return std::stol(name); }
template<> int convert<int>(const std::string& name) { return std::stoi(name); }
class function {
public:
typedef std::function<grid_t(const params_t&)> grid_fn_ty;
struct options_space_t {
typedef std::pair<std::string, std::vector<std::string>> define_t;
std::vector<define_t> defines;
std::vector<size_t> num_warps;
};
struct options_t {
template<class T>
T D(const std::string& name) const {
return convert<T>(defines.at(name));
}
std::map<std::string, std::string> defines;
size_t num_warps;
};
typedef std::function<grid_t(const options_t&)> grid_fn_ty;
private:
class caller {
public:
caller(ir::function *ir, std::shared_ptr<driver::module> program, size_t n_threads);
caller(ir::function *ir, std::shared_ptr<driver::module> program, const options_t& opt_);
void operator()(driver::stream *stream, const std::array<size_t, 3>& grid, const std::vector<arg>& args) const;
const options_t opt() const { return opt_; }
private:
std::shared_ptr<driver::kernel> bin_;
std::shared_ptr<driver::module> parent_;
std::vector<arg_type> param_tys_;
size_t n_threads_;
options_t opt_;
};
private:
typedef std::pair<driver::device*, std::vector<int64_t>> cache_key_t;
typedef std::pair<options, caller> cache_val_t;
private:
triton::lang::translation_unit *make_ast(const std::string &src);
std::unique_ptr<ir::module> make_ir(Parser &parser);
options autotune(Parser &parser, driver::stream *stream, const grid_fn_ty& grid, const std::vector<arg> &args);
std::unique_ptr<driver::module> make_bin(ir::module &function, driver::context *context, const options &opt);
std::unique_ptr<driver::module> make_bin(ir::module &function, driver::context *context, const options_t &opt);
caller autotune(driver::stream *stream, const grid_fn_ty& grid, const std::vector<arg> &args);
public:
function(const std::string& src);
function(const std::string& src, const options_space_t& opt = options_space_t());
void operator()(const std::vector<arg>& args, const std::array<size_t, 3>& grid, driver::stream* stream);
void operator()(const std::vector<arg>& args, const grid_fn_ty& grid, driver::stream *stream);
std::string make_tensorflow_src(const std::vector<size_t> &outputs, const std::string &macro);
private:
TokenSequence ts_;
Parser parser_;
// execution context
ir::context ctx_;
// program representations
std::string src_;
std::map<cache_key_t, cache_val_t> cache_;
std::map<cache_key_t, caller> cache_;
// options
options_space_t opt_space_;
};
}

30
lib/codegen/transform/peephole.cpp
View File

@@ -2,7 +2,7 @@
#include "triton/ir/module.h"
#include "triton/ir/function.h"
#include "triton/codegen/transform/peephole.h"
#include <iostream>
namespace triton {
namespace codegen{
namespace transform{
@@ -145,6 +145,34 @@ bool peephole::rewrite_dot_fp32(ir::dot_inst *dot, ir::builder& builder, bool tr
}
bool peephole::rewrite_dot(ir::instruction *value, ir::builder& builder){
// dot(a, b, 0) + c -> dot(a, b, c)
auto add = dynamic_cast<ir::binary_operator*>(value);
if(add && add->get_op() == ir::binary_op_t::FAdd) {
ir::value *lhs = add->get_operand(0);
ir::value *rhs = add->get_operand(1);
ir::dot_inst *lhs_dot = dynamic_cast<ir::dot_inst*>(lhs);
ir::dot_inst *rhs_dot = dynamic_cast<ir::dot_inst*>(rhs);
if(!lhs_dot && !rhs_dot)
return false;
ir::dot_inst *dot = lhs_dot ? lhs_dot : rhs_dot;
ir::value *other = (dot == lhs) ? rhs : lhs;
ir::value *acc = dot->get_operand(2);
ir::splat_inst *splat = dynamic_cast<ir::splat_inst*>(acc);
ir::constant_fp *_0 = nullptr;
if(splat)
_0 = dynamic_cast<ir::constant_fp*>(splat->get_operand(0));
if(!(_0 && _0->get_value() == 0.0))
return false;
ir::value *a = dot->get_operand(0);
ir::value *b = dot->get_operand(1);
ir::value * new_dot = builder.insert(ir::dot_inst::create(a, b, other,
dot->is_a_trans(), dot->is_b_trans(),
dot->get_name()));
add->replace_all_uses_with(new_dot);
return true;
}
// dot(a, b, c)
auto dot = dynamic_cast<ir::dot_inst*>(value);
if(!dot)
return false;

9
lib/driver/module.cpp
View File

@@ -106,10 +106,10 @@ void module::compile_llvm_module(llvm::Module* module, const std::string& triple
file_type_t ft) {
init_llvm();
// debug
llvm::legacy::PassManager pm;
pm.add(llvm::createPrintModulePass(llvm::outs()));
pm.add(llvm::createVerifierPass());
pm.run(*module);
// llvm::legacy::PassManager pm;
// pm.add(llvm::createPrintModulePass(llvm::outs()));
// pm.add(llvm::createVerifierPass());
// pm.run(*module);
// create machine
module->setTargetTriple(triple);
std::string error;
@@ -221,7 +221,6 @@ ocl_module::ocl_module(driver::context * context, llvm::Module* src): module(con
catch(...){
char log[2048];
dispatch::clGetProgramBuildInfo(*cl_, *context->device()->cl(), CL_PROGRAM_BUILD_LOG, 1024, log, NULL);
std::cout << log << std::endl;
throw;
}
}

115
lib/runtime/function.cpp
View File

@@ -17,7 +17,6 @@
#include "triton/ir/function.h"
#include "triton/tools/bench.hpp"
#include "llvm/IR/Module.h"
#include "triton/ir/print.h"
typedef struct yy_buffer_state * YY_BUFFER_STATE;
@@ -91,8 +90,8 @@ arg_type convert(ir::type *ty) {
throw std::runtime_error("unknown type");
}
function::caller::caller(ir::function *ir, std::shared_ptr<driver::module> parent, size_t n_threads)
: bin_(driver::kernel::create(&*parent, ir->get_name().c_str())), parent_(parent), n_threads_(n_threads) {
function::caller::caller(ir::function *ir, std::shared_ptr<driver::module> parent, const options_t& opt)
: bin_(driver::kernel::create(&*parent, ir->get_name().c_str())), parent_(parent), opt_(opt) {
// extract signature
ir::function_type* ty = ir->get_fn_type();
for(size_t i = 0; i < ty->get_num_params(); i++)
@@ -113,12 +112,10 @@ void function::caller::operator ()(driver::stream *stream, const std::array<size
else
bin_->setArg(i, size_of(ty), arg_i.data());
}
stream->enqueue(&*bin_, grid, {n_threads_, 1, 1});
stream->enqueue(&*bin_, grid, {opt_.num_warps * 32, 1, 1});
}
std::unique_ptr<ir::module> function::make_ir(Parser& parser) {
// create Triton-IR from AST
ir::module* module = new ir::module("", ctx_);
@@ -127,59 +124,59 @@ std::unique_ptr<ir::module> function::make_ir(Parser& parser) {
return std::unique_ptr<ir::module>(module);
}
options function::autotune(Parser& parser, driver::stream* stream, const grid_fn_ty& grid_fn, const std::vector<arg>& args) {
std::unique_ptr<ir::module> ir = make_ir(parser);
// extract tunable values
std::vector<std::pair<std::string, ir::metaparameter*>> values;
for(auto it: ir->globals())
if(auto *mp = dynamic_cast<ir::metaparameter*>(it.second))
values.push_back({it.first, mp});
// extract search space
std::vector<std::vector<unsigned>> space;
space.push_back({1, 2, 4, 8}); // num warps
for(auto it: values)
space.push_back(it.second->get_space());
function::caller function::autotune(driver::stream* stream, const grid_fn_ty& grid_fn,
const std::vector<arg>& args) {
// all tuning parameters are strings
std::vector<std::string> num_warps;
for(size_t i: opt_space_.num_warps)
num_warps.push_back(std::to_string(i));
std::vector<std::vector<std::string>> space;
space.push_back(num_warps);
for(const auto& i: opt_space_.defines)
space.push_back(i.second);
// exhaustive search
struct profile_t{
double ts;
std::vector<unsigned> params;
};
profile_t best = { INFINITY, {} };
std::function<void(std::vector<unsigned>)> benchmark =
[&](std::vector<unsigned> params) {
// options
options opt;
double best_ts = INFINITY;
std::unique_ptr<caller> ret;
auto benchmark = [&](std::vector<std::string> params) {
// extract options
options_t opt;
unsigned i = 0;
opt.num_warps = params[i++];
for(auto it: values)
opt.params[it.first] = params[i++];
// make binary
opt.num_warps = std::stoi(params[i++]);
for(auto it: opt_space_.defines)
opt.defines[it.first] = params[i++];
// pre-process
TokenSequence tokens;
Preprocessor cpp(&src_, true);
for(auto it: opt_space_.defines)
cpp.AddMacro(it.first, &opt.defines.at(it.first));
cpp.Process(tokens);
// parse
Parser parser(tokens);
parser.Parse();
// triton-ir code-gen
auto ir = make_ir(parser);
// binary code-gen
auto bin = make_bin(*ir, stream->context(), opt);
// benchmark
ir::function *tmp = ir->get_function_list()[0];
caller fn(tmp, std::move(bin), opt.num_warps * 32);
double ts = tools::bench([&]() { fn(stream, grid_fn(opt.params), args); }, stream);
if(ts < best.ts)
best = {ts, params};
caller call(tmp, std::move(bin), opt);
double ts = tools::bench([&]() { call(stream, grid_fn(opt), args); }, stream);
// save best
if(ts < best_ts)
ret.reset(new caller(call));
};
_parallel_loop_nest<unsigned>(space, benchmark, 1);
// populate options
unsigned current = 0;
options opt;
opt.num_warps = best.params[current++];
for(auto it: values)
opt.params[it.first] = best.params[current++];
return opt;
_parallel_loop_nest<std::string>(space, benchmark, 1);
return *ret;
}
std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::context *context, const options& opt) {
std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::context *context, const options_t& opt) {
std::unique_ptr<codegen::target> target = context->device()->make_target();
// update metaparameter values
for(auto x: opt.params)
if(auto* mp = dynamic_cast<ir::metaparameter*>(module.globals().at(x.first)))
mp->set_value(x.second);
// create passes
codegen::analysis::grids tune(opt.num_warps);
codegen::analysis::shmem::info shmem_info;
@@ -210,8 +207,6 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
// generate llvm code
llvm::LLVMContext ctx;
std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx));
ir::print(module, std::cout);
exit(EXIT_FAILURE);
selection.run(module, *llvm);
// return binary
std::unique_ptr<driver::module> res(driver::module::create(context, llvm.get()));
@@ -219,11 +214,8 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
}
function::function(const std::string &src): parser_(ts_), src_(src){
Preprocessor cpp(&src_, true);
cpp.Process(ts_);
ts_.Print();
parser_.Parse();
function::function(const std::string &src, const options_space_t& opt): src_(src), opt_space_(opt) {
}
void function::operator()(const std::vector<arg>& args, const grid_fn_ty& grid_fn, driver::stream *stream) {
@@ -244,21 +236,18 @@ void function::operator()(const std::vector<arg>& args, const grid_fn_ty& grid_f
/* find existing configuration */
auto it = cache_.find(key);
if(it != cache_.end()){
it->second.second(stream, grid_fn(it->second.first.params), args);
it->second(stream, grid_fn(it->second.opt()), args);
return;
}
/* re-tune and re-compile */
options opt = autotune(parser_, stream, grid_fn, args);
std::unique_ptr<ir::module> ir = make_ir(parser_);
std::unique_ptr<driver::module> bin = make_bin(*ir, stream->context(), opt);
ir::function* fn = ir->get_function_list().front();
const caller& run = cache_.insert({key, cache_val_t{opt, caller(fn, std::move(bin), opt.num_warps*32)}}).first->second.second;
run(stream, grid_fn(opt.params), args);
caller call = autotune(stream, grid_fn, args);
cache_.insert({key, call});
}
void function::operator()(const std::vector<arg>& args, const grid_t& grid, driver::stream *stream) {
return this->operator()(args, [&grid](const params_t&){ return grid; }, stream);
return this->operator()(args, [&grid](const options_t&){ return grid; }, stream);
}
}