diff --git a/examples/cpp/dot.cpp b/examples/cpp/dot.cpp index e93a09994..b56f6f6e1 100644 --- a/examples/cpp/dot.cpp +++ b/examples/cpp/dot.cpp @@ -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 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}, diff --git a/include/triton/codegen/selection/selection.h b/include/triton/codegen/selection/selection.h index d17c0607b..785e32179 100644 --- a/include/triton/codegen/selection/selection.h +++ b/include/triton/codegen/selection/selection.h @@ -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{ diff --git a/include/triton/runtime/function.h b/include/triton/runtime/function.h index e43301bfc..2cbd65fd4 100644 --- a/include/triton/runtime/function.h +++ b/include/triton/runtime/function.h @@ -97,6 +97,7 @@ public: function(const std::string& src); void operator()(const std::vector& args, const std::array& grid, driver::stream* stream); void operator()(const std::vector& args, const grid_fn_ty& grid, driver::stream *stream); + std::string make_tensorflow_src(const std::vector &outputs, const std::string ¯o); private: // execution context diff --git a/lib/runtime/function.cpp b/lib/runtime/function.cpp index fde0c7972..24e66397b 100644 --- a/lib/runtime/function.cpp +++ b/lib/runtime/function.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include "triton/codegen/selection/selection.h" #include "triton/runtime/function.h" @@ -259,6 +260,173 @@ void function::operator()(const std::vector& 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& outputs, const std::string& macro) { + std::unique_ptr 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 arg_names; + for(ir::argument *arg: fn->args()) + arg_names.push_back(arg->get_name()); + // cached int to str + std::vector 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 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 tf_tys; + std::transform(fn_ty->params_begin(), fn_ty->params_end(), std::back_inserter(tf_tys), to_tf_ty); + std::vector 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(); + 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; +} + + }