Deprecation of Triton-C and Replacement by decorated Python functions (#86)
This PR implements a major overhaul of the frontend for Triton, and replaces Triton-C by a pure Python API in which kernels are defined as @triton.jit decorated functions. The documentation and tutorials have also been updated to accommodate these changes. See documentations for more information on the new API
This commit is contained in:
Philippe Tillet
committed by
Philippe Tillet
Philippe Tillet
parent
1fdb465b71
commit
39f4730305
@@ -1,5 +1,13 @@
|
||||
#include "triton/driver/stream.h"
|
||||
#include "triton/runtime/function.h"
|
||||
#include "triton/codegen/pass.h"
|
||||
#include "triton/driver/kernel.h"
|
||||
#include "triton/driver/module.h"
|
||||
#include "triton/driver/stream.h"
|
||||
#include "triton/ir/builder.h"
|
||||
#include "triton/ir/dispatch.h"
|
||||
#include "triton/ir/enums.h"
|
||||
#include "triton/ir/function.h"
|
||||
#include "triton/ir/module.h"
|
||||
#include <optional>
|
||||
#include <pybind11/buffer_info.h>
|
||||
#include <pybind11/functional.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
@@ -8,78 +16,9 @@
|
||||
#include <string>
|
||||
|
||||
namespace py = pybind11;
|
||||
|
||||
using namespace triton;
|
||||
namespace rt = triton::runtime;
|
||||
namespace ir = triton::ir;
|
||||
namespace drv = triton::driver;
|
||||
|
||||
/*****************************************************************************/
|
||||
/* Python bindings for triton::tools */
|
||||
/*****************************************************************************/
|
||||
|
||||
/*!
|
||||
@brief Function for extracting kernels out of a given source-string
|
||||
|
||||
This can be important to enable pre-processor macros (or tunable parameters) that should only
|
||||
be defined within the scope of a single kernel function
|
||||
*/
|
||||
std::string extract_kernels(const std::string &str, const std::vector<std::string> &names) {
|
||||
if (names.empty())
|
||||
return str;
|
||||
// search for all regex matches of kernel_regex in str
|
||||
std::smatch matches;
|
||||
std::regex regex(" *__global__ +void +([_a-zA-Z][_a-zA-Z0-9]{0,30})");
|
||||
std::sregex_iterator it(str.begin(), str.end(), regex);
|
||||
std::sregex_iterator end;
|
||||
std::vector<std::tuple<std::string, int, int>> kernels;
|
||||
for (; it != end; ++it) {
|
||||
int pos = it->position();
|
||||
int len = it->length();
|
||||
std::string name = it->str(1);
|
||||
kernels.push_back(std::make_tuple(name, pos, len));
|
||||
}
|
||||
// check that all the kernels provided actually exist
|
||||
for (const std::string &name : names) {
|
||||
auto pred = [&name](const std::tuple<std::string, int, int> &t) { return std::get<0>(t) == name; };
|
||||
bool found = std::any_of(kernels.begin(), kernels.end(), pred);
|
||||
if (!found)
|
||||
throw std::runtime_error("Unable to find kernel `" + name + "` in provided source code:\n" + str);
|
||||
}
|
||||
// simple parsing logic to extract the declaration and body of each specified kernel
|
||||
std::string ret;
|
||||
for (const auto &k : kernels) {
|
||||
std::string name;
|
||||
int pos, len;
|
||||
std::tie(name, pos, len) = k;
|
||||
if (std::find(names.begin(), names.end(), name) == names.end())
|
||||
continue;
|
||||
std::string def = str.substr(pos, str.size() - pos);
|
||||
// skip over declaration
|
||||
// by finding matching ')' for first '('
|
||||
int count = 1;
|
||||
pos = def.find('(');
|
||||
while (!(def[pos++] == ')' && count == 0) && pos < def.size()) {
|
||||
count += def[pos] == '(';
|
||||
count -= def[pos] == ')';
|
||||
}
|
||||
// skip over definition
|
||||
// by finding matching '{' for first '}'
|
||||
count = 1;
|
||||
pos = def.find('{', pos);
|
||||
while (!(def[pos++] == '}' && count == 0) && pos < def.size()) {
|
||||
count += def[pos] == '{';
|
||||
count -= def[pos] == '}';
|
||||
}
|
||||
ret += def.substr(0, pos);
|
||||
ret += '\n';
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
void init_triton_tools(py::module &&m) {
|
||||
m.def("extract_kernels", &extract_kernels);
|
||||
}
|
||||
|
||||
/*****************************************************************************/
|
||||
/* Python bindings for triton::driver */
|
||||
/*****************************************************************************/
|
||||
@@ -88,14 +27,14 @@ void init_triton_driver(py::module &&m) {
|
||||
// base device
|
||||
py::class_<drv::device>(m, "device");
|
||||
// cuda device
|
||||
py::class_<drv::cu_device, driver::device>(m, "cu_device")
|
||||
py::class_<drv::cu_device, drv::device>(m, "cu_device")
|
||||
.def(py::init([](int dev_id, bool take_ownership) {
|
||||
CUdevice handle;
|
||||
drv::dispatch::cuDeviceGet(&handle, dev_id);
|
||||
return new drv::cu_device(handle, take_ownership);
|
||||
}));
|
||||
// host device
|
||||
py::class_<drv::host_device, driver::device>(m, "host_device")
|
||||
py::class_<drv::host_device, drv::device>(m, "host_device")
|
||||
.def(py::init<>());
|
||||
|
||||
// base stream
|
||||
@@ -108,54 +47,236 @@ void init_triton_driver(py::module &&m) {
|
||||
// py doesn't support opaque pointer (e.g., CUstream) so
|
||||
// we assume it has been converted to uint64_t
|
||||
.def(py::init([](uint64_t handle, bool take_ownership) {
|
||||
return std::unique_ptr<driver::cu_stream>(new driver::cu_stream((CUstream)handle, take_ownership));
|
||||
}));
|
||||
return std::unique_ptr<drv::cu_stream>(new drv::cu_stream((CUstream)handle, take_ownership));
|
||||
}))
|
||||
.def("enqueue", [](drv::cu_stream *self, drv::kernel *kernel,
|
||||
size_t grid_0, size_t grid_1, size_t grid_2,
|
||||
size_t block_0, size_t block_1, size_t block_2,
|
||||
const std::string &args,
|
||||
size_t shared_mem) {
|
||||
return self->enqueue(kernel, {grid_0, grid_1, grid_2}, {block_0, block_1, block_2},
|
||||
(void *)args.data(), args.size(), shared_mem);
|
||||
});
|
||||
|
||||
py::class_<drv::module>(m, "module");
|
||||
//py::class_<drv::cu_module, drv::module>(m, "cu_module");
|
||||
|
||||
py::class_<drv::kernel>(m, "kernel");
|
||||
}
|
||||
|
||||
/*****************************************************************************/
|
||||
/* Python bindings for triton::runtime */
|
||||
/* Python bindings for triton::codegen */
|
||||
/*****************************************************************************/
|
||||
void init_triton_runtime(py::module &&m) {
|
||||
// argument type
|
||||
py::enum_<rt::arg_type>(m, "arg_type")
|
||||
.value("int1", rt::INT1_T)
|
||||
.value("int8", rt::INT8_T)
|
||||
.value("int16", rt::INT16_T)
|
||||
.value("int32", rt::INT32_T)
|
||||
.value("int64", rt::INT64_T)
|
||||
.value("half", rt::HALF_T)
|
||||
.value("float", rt::FLOAT_T)
|
||||
.value("double", rt::DOUBLE_T)
|
||||
.value("buffer", rt::BUFFER_T);
|
||||
// compilation options
|
||||
py::class_<rt::options_t>(m, "options", py::dynamic_attr())
|
||||
.def(py::init<>())
|
||||
.def_readwrite("defines", &rt::options_t::defines)
|
||||
.def_readwrite("num_warps", &rt::options_t::num_warps)
|
||||
.def("__getattr__", [](rt::options_t *opt, const std::string &name) {
|
||||
return opt->D<int>(name);
|
||||
});
|
||||
// kernel
|
||||
py::class_<rt::kernel>(m, "kernel")
|
||||
.def("__call__", &rt::kernel::operator())
|
||||
.def_readonly("opt", &rt::kernel::opt)
|
||||
.def("asm", &rt::kernel::get_asm);
|
||||
// tune conf
|
||||
py::class_<rt::config>(m, "config")
|
||||
.def(py::init<std::map<std::string, std::string>, int>(),
|
||||
py::arg("defines") = std::map<std::string, std::string>(),
|
||||
py::arg("num_warps"));
|
||||
|
||||
// function
|
||||
py::class_<rt::function>(m, "function")
|
||||
.def(py::init<const std::string &, const rt::options_t &, driver::device *, const std::vector<rt::config> &, const std::vector<std::string> &>())
|
||||
.def("autotune", &rt::function::autotune, py::return_value_policy::reference_internal)
|
||||
.def("signature", &rt::function::get_signature);
|
||||
void init_triton_codegen(py::module &&m) {
|
||||
m.def(
|
||||
"add_passes_to_emit_bin", [](ir::module &ir, drv::device *dev, int num_warps) {
|
||||
drv::module *mod;
|
||||
drv::kernel *ker;
|
||||
size_t shared_mem;
|
||||
triton::codegen::add_passes_to_emit_bin(ir, dev, num_warps, mod, ker, shared_mem);
|
||||
return std::make_tuple(mod, ker, shared_mem);
|
||||
},
|
||||
py::return_value_policy::take_ownership);
|
||||
}
|
||||
|
||||
/*****************************************************************************/
|
||||
/* User-facing language features */
|
||||
/*****************************************************************************/
|
||||
|
||||
void init_triton_frontend(py::module &&m) {
|
||||
using ret = py::return_value_policy;
|
||||
|
||||
// programming model
|
||||
m.def("program_id", &ir::dispatch::program_id, ret::reference);
|
||||
m.def("num_programs", &ir::dispatch::num_programs, ret::reference);
|
||||
// binary
|
||||
m.def("add", &ir::dispatch::add, ret::reference);
|
||||
m.def("sub", &ir::dispatch::sub, ret::reference);
|
||||
m.def("mul", &ir::dispatch::mul, ret::reference);
|
||||
m.def("truediv", &ir::dispatch::truediv, ret::reference);
|
||||
m.def("floordiv", &ir::dispatch::floordiv, ret::reference);
|
||||
m.def("mod", &ir::dispatch::mod, ret::reference);
|
||||
m.def("and_", &ir::dispatch::and_, ret::reference);
|
||||
m.def("or_", &ir::dispatch::or_, ret::reference);
|
||||
m.def("xor_", &ir::dispatch::xor_, ret::reference);
|
||||
m.def("lshr", &ir::dispatch::lshr, ret::reference);
|
||||
m.def("shl", &ir::dispatch::shl, ret::reference);
|
||||
// unary
|
||||
m.def("plus", &ir::dispatch::plus, ret::reference);
|
||||
m.def("minus", &ir::dispatch::minus, ret::reference);
|
||||
m.def("invert", &ir::dispatch::invert, ret::reference);
|
||||
// comparison
|
||||
m.def("greater_than", &ir::dispatch::greater_than, ret::reference);
|
||||
m.def("greater_equal", &ir::dispatch::greater_equal, ret::reference);
|
||||
m.def("less_than", &ir::dispatch::less_than, ret::reference);
|
||||
m.def("less_equal", &ir::dispatch::less_equal, ret::reference);
|
||||
m.def("equal", &ir::dispatch::equal, ret::reference);
|
||||
m.def("not_equal", &ir::dispatch::not_equal, ret::reference);
|
||||
// block creation
|
||||
m.def("arange", &ir::dispatch::arange, ret::reference);
|
||||
m.def("zeros", &ir::dispatch::zeros, ret::reference);
|
||||
// type manipuatation
|
||||
m.def("reshape", &ir::dispatch::reshape, ret::reference);
|
||||
typedef std::tuple<ir::value *, ir::value *> (*broadcast_ty)(ir::value *, ir::value *, ir::builder *);
|
||||
typedef ir::value *(*broadcast_to_ty)(ir::value *, ir::type::block_shapes_t, ir::builder *);
|
||||
m.def("broadcast", (broadcast_ty)(&ir::dispatch::broadcast), ret::reference);
|
||||
m.def("broadcast_to", (broadcast_to_ty)(&ir::dispatch::broadcast), ret::reference);
|
||||
m.def("cast", &ir::dispatch::cast, ret::reference);
|
||||
// memory
|
||||
m.def("load", &ir::dispatch::load, ret::reference);
|
||||
m.def("store", &ir::dispatch::store, ret::reference);
|
||||
m.def("atomic_cas", &ir::dispatch::atomic_cas, ret::reference);
|
||||
m.def("atomic_xchg", &ir::dispatch::atomic_xchg, ret::reference);
|
||||
// linear algebra
|
||||
m.def("dot", &ir::dispatch::dot, ret::reference);
|
||||
// indexing
|
||||
m.def("where", &ir::dispatch::where, ret::reference);
|
||||
// reduction
|
||||
m.def("min", &ir::dispatch::min, ret::reference);
|
||||
m.def("max", &ir::dispatch::max, ret::reference);
|
||||
m.def("sum", &ir::dispatch::sum, ret::reference);
|
||||
// math
|
||||
m.def("exp", &ir::dispatch::exp, ret::reference);
|
||||
m.def("log", &ir::dispatch::log, ret::reference);
|
||||
m.def("sqrt", &ir::dispatch::sqrt, ret::reference);
|
||||
// internal (debugging only)
|
||||
m.def("multiple_of", &ir::dispatch::multiple_of, ret::reference);
|
||||
m.def("debug_barrier", &ir::dispatch::debug_barrier, ret::reference);
|
||||
}
|
||||
|
||||
/*****************************************************************************/
|
||||
/* Python bindings for triton::ir */
|
||||
/*****************************************************************************/
|
||||
|
||||
void init_triton_ir(py::module &&m) {
|
||||
using ret = py::return_value_policy;
|
||||
using namespace pybind11::literals;
|
||||
|
||||
py::class_<ir::context>(m, "context")
|
||||
.def(py::init<>());
|
||||
|
||||
auto value = py::class_<ir::value>(m, "value");
|
||||
value.def_property("name", &ir::value::get_name, &ir::value::set_name);
|
||||
value.def_property_readonly("type", &ir::value::get_type);
|
||||
|
||||
py::class_<ir::user, ir::value>(m, "user");
|
||||
|
||||
py::class_<ir::constant, ir::user>(m, "constant");
|
||||
|
||||
py::class_<ir::undef_value, ir::constant>(m, "undef")
|
||||
.def("get", &ir::undef_value::get, ret::reference);
|
||||
|
||||
py::class_<ir::constant_int, ir::constant>(m, "constant_int")
|
||||
.def_property_readonly("value", &ir::constant_int::get_value)
|
||||
.def("__int__", [](ir::constant_int *self) { return self->get_value(); });
|
||||
|
||||
py::class_<ir::constant_fp, ir::constant>(m, "constant_float")
|
||||
.def_property_readonly("value", &ir::constant_fp::get_value);
|
||||
|
||||
py::class_<ir::type>(m, "type")
|
||||
.def("is_ptr", &ir::type::is_pointer_ty)
|
||||
.def("is_int", static_cast<bool (ir::type::*)() const>(&ir::type::is_integer_ty))
|
||||
.def("is_floating", &ir::type::is_floating_point_ty)
|
||||
.def("is_block", &ir::type::is_block_ty)
|
||||
.def("make_ptr", &ir::pointer_type::get, ret::reference)
|
||||
.def("make_function", &ir::function_type::get, ret::reference)
|
||||
.def("make_block", &ir::block_type::get, ret::reference)
|
||||
.def("get_void", &ir::type::get_void_ty, ret::reference)
|
||||
.def("get_fp16", &ir::type::get_half_ty, ret::reference)
|
||||
.def("get_fp32", &ir::type::get_float_ty, ret::reference)
|
||||
.def("get_fp64", &ir::type::get_double_ty, ret::reference)
|
||||
.def("get_int1", &ir::type::get_int1_ty, ret::reference)
|
||||
.def("get_int8", &ir::type::get_int8_ty, ret::reference)
|
||||
.def("get_int16", &ir::type::get_int16_ty, ret::reference)
|
||||
.def("get_int32", &ir::type::get_int32_ty, ret::reference)
|
||||
.def("get_int64", &ir::type::get_int64_ty, ret::reference)
|
||||
|
||||
.def("is_void", &ir::type::is_void_ty)
|
||||
.def("is_fp16", &ir::type::is_half_ty)
|
||||
.def("is_fp32", &ir::type::is_float_ty)
|
||||
.def("is_fp64", &ir::type::is_double_ty)
|
||||
.def("is_int1", [](ir::type *self) { return self->is_integer_ty(1); })
|
||||
.def("is_int8", [](ir::type *self) { return self->is_integer_ty(8); })
|
||||
.def("is_int16", [](ir::type *self) { return self->is_integer_ty(16); })
|
||||
.def("is_int32", [](ir::type *self) { return self->is_integer_ty(32); })
|
||||
.def("is_int64", [](ir::type *self) { return self->is_integer_ty(64); })
|
||||
|
||||
.def_property_readonly("fp_mantissa_width", &ir::type::get_fp_mantissa_width)
|
||||
.def_property_readonly("scalar", &ir::type::get_scalar_ty)
|
||||
.def_property_readonly("context", &ir::type::get_context, ret::reference);
|
||||
|
||||
py::class_<ir::pointer_type, ir::type>(m, "pointer_type")
|
||||
.def_property_readonly("element", &ir::pointer_type::get_element_ty, ret::reference);
|
||||
|
||||
py::class_<ir::function_type, ir::type>(m, "function_type");
|
||||
py::class_<ir::integer_type, ir::type>(m, "integer_type");
|
||||
py::class_<ir::block_type, ir::type>(m, "block_type")
|
||||
.def_property_readonly("shape", &ir::block_type::get_shapes)
|
||||
.def_property_readonly("numel", &ir::type::get_tile_num_elements);
|
||||
|
||||
py::class_<ir::scope>(m, "scope")
|
||||
.def(py::init<>())
|
||||
.def_property_readonly("values", &ir::scope::get_values)
|
||||
.def("set_type", &ir::scope::set_type);
|
||||
|
||||
py::class_<ir::module>(m, "module")
|
||||
.def(py::init<std::string, ir::builder &>())
|
||||
.def("get_or_insert_function", &ir::module::get_or_insert_function, ret::reference)
|
||||
.def("add_new_scope", &ir::module::add_new_scope, ret::reference)
|
||||
.def("seal_block", &ir::module::seal_block)
|
||||
.def("set_value", (void (ir::module::*)(const std::string &, ir::value *)) & ir::module::set_value)
|
||||
.def("get_value", (ir::value * (ir::module::*)(const std::string &)) & ir::module::get_value, ret::reference)
|
||||
.def("pop_scope", &ir::module::pop_scope)
|
||||
.def_property_readonly("scope", &ir::module::get_scope, ret::reference)
|
||||
.def_property_readonly("builder", &ir::module::get_builder, ret::reference);
|
||||
|
||||
using eattr = ir::attribute_kind_t;
|
||||
py::enum_<eattr>(m, "attribute_kind")
|
||||
.value("readonly", eattr::readonly)
|
||||
.value("writeonly", eattr::writeonly)
|
||||
.value("noalias", eattr::noalias)
|
||||
.value("aligned", eattr::aligned)
|
||||
.value("multiple_of", eattr::multiple_of)
|
||||
.value("retune", eattr::retune)
|
||||
.value("not_implemented", eattr::not_implemented);
|
||||
|
||||
py::class_<ir::attribute>(m, "attribute")
|
||||
.def(py::init<eattr, int>());
|
||||
|
||||
py::class_<ir::function>(m, "function")
|
||||
.def_property_readonly("args", &ir::function::args)
|
||||
.def_property_readonly("attrs", &ir::function::attrs)
|
||||
.def("add_attr", &ir::function::add_attr);
|
||||
|
||||
py::class_<ir::argument, ir::value>(m, "argument");
|
||||
|
||||
py::class_<ir::basic_block, ir::value>(m, "basic_block")
|
||||
.def("create", &ir::basic_block::create, ret::reference)
|
||||
.def_property_readonly("parent", &ir::basic_block::get_parent, ret::reference);
|
||||
|
||||
py::class_<ir::builder>(m, "builder", py::dynamic_attr())
|
||||
.def(py::init<ir::context &>())
|
||||
// getters
|
||||
.def_property_readonly("context", &ir::builder::get_context, ret::reference)
|
||||
// control flow
|
||||
.def("br", &ir::builder::create_br, ret::reference)
|
||||
.def("cond_br", &ir::builder::create_cond_br, ret::reference)
|
||||
.def("ret_void", &ir::builder::create_ret_void, ret::reference)
|
||||
.def("get_insert_block", &ir::builder::get_insert_block, ret::reference)
|
||||
.def("set_insert_block", (void (ir::builder::*)(ir::basic_block *)) & ir::builder::set_insert_point)
|
||||
// constants
|
||||
.def("get_int1", &ir::builder::get_int1, ret::reference)
|
||||
.def("get_int32", &ir::builder::get_int32, ret::reference)
|
||||
.def("get_float16", &ir::builder::get_float16, ret::reference)
|
||||
.def("get_float32", &ir::builder::get_float32, ret::reference)
|
||||
.def("get_range", &ir::builder::get_range, ret::reference);
|
||||
}
|
||||
|
||||
void init_triton(py::module &m) {
|
||||
py::module subm = m.def_submodule("triton");
|
||||
init_triton_codegen(std::move(subm.def_submodule("code_gen")));
|
||||
init_triton_driver(std::move(subm.def_submodule("driver")));
|
||||
init_triton_runtime(std::move(subm.def_submodule("runtime")));
|
||||
init_triton_tools(std::move(subm.def_submodule("tools")));
|
||||
init_triton_ir(std::move(subm.def_submodule("ir")));
|
||||
init_triton_frontend(std::move(subm.def_submodule("frontend")));
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user