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

added tensorflow code generator

Browse Source
This commit is contained in:
Philippe Tillet
2019-08-15 15:59:53 -07:00
parent 38a8b0ab19
commit 3ece461ce2
4 changed files with 174 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/cpp/dot.cpp
View File

@@ -154,6 +154,9 @@ perf_t do_bench(drv::stream* stream, bool AT, bool BT, int32_t M, int32_t N, int
stream->synchronize();
// run
rt::function function(src(AT, BT, ty, ty, ty, 8, 8));
std::cout << function.make_tensorflow_src({2}, "(M + #TM - 1)/#TM, (N + #TN - 1)/#TN, 1") << std::endl;
exit(EXIT_FAILURE);
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}; };
@@ -202,7 +205,7 @@ int main() {
// shapes to benchmark
std::vector<config_t> configs = {
// {false, false, 8192, 512, 512},
{false, true, 8192, 8192, 8192}
{false, true, 128, 128, 128}
// {false, true, 128, 128, 128},
// {false, false, 128, 128, 128},
// {true, false, 128, 128, 128},

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

@@ -1,7 +1,6 @@
#ifndef TDL_INCLUDE_CODEGEN_SELECTION_H
#define TDL_INCLUDE_CODEGEN_SELECTION_H
#include "llvm/IR/Module.h"
#include "llvm/IR/IRBuilder.h"
#include "triton/ir/context.h"
#include "triton/ir/module.h"
@@ -16,6 +15,7 @@ namespace llvm{
class Instruction;
class Constant;
class LLVMContext;
class Module;
}
namespace triton{

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

@@ -97,6 +97,7 @@ public:
function(const std::string& src);
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:
// execution context

168
lib/runtime/function.cpp
View File

@@ -1,5 +1,6 @@
#include <string>
#include <mutex>
#include <regex>
#include <functional>
#include "triton/codegen/selection/selection.h"
#include "triton/runtime/function.h"
@@ -259,6 +260,173 @@ void function::operator()(const std::vector<arg>& args, const grid_t& grid, driv
return this->operator()(args, [&grid](const params_t&){ return grid; }, stream);
}
std::string to_tf_ty(ir::type *ty) {
if(ty->is_integer_ty(1))
return "bool";
if(ty->is_integer_ty(8))
return "int8";
if(ty->is_integer_ty(16))
return "int16";
if(ty->is_integer_ty(32))
return "int32";
if(ty->is_integer_ty(64))
return "int64";
if(ty->is_half_ty())
return "float16";
if(ty->is_float_ty())
return "float32";
if(ty->is_double_ty())
return "float64";
if(ty->is_pointer_ty())
return "Tensor";
throw std::runtime_error("unknown type");
}
std::string ref_to_tf_ty(ir::type *ty) {
std::string res = to_tf_ty(ty);
if(ty->is_pointer_ty())
res = "const " + res + "&";
return res;
}
std::string function::make_tensorflow_src(const std::vector<size_t>& outputs, const std::string& macro) {
std::unique_ptr<ir::module> ir = make_ir(ast_);
// extract function signature
ir::function* fn = ir->get_function_list().front();
ir::function_type* fn_ty = fn->get_fn_type();
// numberof arguments
size_t n_args = fn_ty->get_num_params();
size_t n_outputs = outputs.size();
// extract function name
std::string name = fn->get_name();
std::string classname = name + "Op";
// extract argument name
std::vector<std::string> arg_names;
for(ir::argument *arg: fn->args())
arg_names.push_back(arg->get_name());
// cached int to str
std::vector<std::string> str_i;
for(size_t i = 0; i < fn_ty->get_num_params(); i++)
str_i.push_back(std::to_string(i));
// index of tensors
std::vector<size_t> ptr_idx;
for(unsigned i = 0; i < fn_ty->get_num_params(); i++)
if(fn_ty->get_param_ty(i)->is_pointer_ty())
ptr_idx.push_back(i);
// extract tensorflow types
std::vector<std::string> tf_tys;
std::transform(fn_ty->params_begin(), fn_ty->params_end(), std::back_inserter(tf_tys), to_tf_ty);
std::vector<std::string> tf_cref_tys;
std::transform(fn_ty->params_begin(), fn_ty->params_end(), std::back_inserter(tf_cref_tys), ref_to_tf_ty);
std::ostringstream oss;
std::string result = R"(
#include "triton/driver/buffer.h"
#include "triton/driver/backend.h"
#include "triton/driver/stream.h"
#include "triton/runtime/function.h"
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"
#include "tensorflow/core/util/padding.h"
#include "tensorflow/core/util/tensor_format.h"
#include "tensorflow/core/framework/common_shape_fns.h"
using namespace tensorflow;
using GPUDevice = Eigen::GpuDevice;
namespace rt = triton::runtime;
namespace drv = triton::driver;
std::string src = R"TTKERNSRC( )" + src_ + ")TTKERNSRC\";" + R"(
class )" + classname + R"(: public OpKernel {
public:
explicit )" + classname + R"((OpKernelConstruction* context)
: OpKernel(context), fn_(src) { }
void Compute(OpKernelContext* context){
// get device/stream
GPUDevice device = context->eigen_device<GPUDevice>();
drv::cu_stream sstream(device.stream(), false);
drv::context* ctx = sstream.context();
drv::stream* stream = &sstream;
// extract inputs)";
for(unsigned i = 0; i < n_args; i++){
std::string suffix = "";
std::string ty = tf_cref_tys[i];
if(!fn_ty->get_param_ty(i)->is_pointer_ty())
suffix = ".scalar<" + ty + ">()()";
result += R"(
)" + ty + " " + arg_names[i] + " = context->input(" + str_i[i] + ")" + suffix + ";";
}
result += R"(
// extract outputs)";
for(unsigned i = 0; i < n_outputs; i++)
result += R"(
context->set_output()" + str_i[i] + ", " + arg_names[outputs[i]] + ");";
result += R"(
// wrap tensors)";
for(size_t i: ptr_idx)
result += R"(
drv::cu_buffer cu_)" + arg_names[i] + "(ctx, " + arg_names[i] + ".tensor_data().size(), (CUdeviceptr)" + arg_names[i] + R"(.tensor_data().data(), false);)";
std::regex regex("#([a-zA-Z]([a-zA-Z]|[0-9])*)");
std::string grid_str = std::regex_replace(macro, regex, "x.at(\"$1\")");
result += R"(
// create launch grid;
auto grid = [&](const rt::params_t& x) { return rt::grid_t{)" + grid_str + R"(}; };)";
result += R"(
// execute function
fn_({
)";
for(unsigned i = 0; i < n_args; i++){
std::string arg = arg_names[i];
if(fn_ty->get_param_ty(i)->is_pointer_ty())
arg = "&cu_" + arg;
if(i > 0)
result += ", ";
result += arg;
}
result += R"(
}, grid, stream);
}
private:
rt::function fn_;
};
REGISTER_KERNEL_BUILDER(Name(")" + name + "\").Device(DEVICE_GPU), " + classname + R"();
REGISTER_OP(")" + name + "\")\n";
for(size_t i = 0; i < tf_tys.size(); i++){
bool is_output = std::find(outputs.begin(), outputs.end(), i) != outputs.end();
std::string mode = is_output ? "Output" : "Input" ;
result += " ." + mode + "(\"" + arg_names[i] + ": " + tf_tys[i] + "\")\n";
}
result += ";\n";
return result;
}
}