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

Fix OpBuilder

Browse Source
This commit is contained in:
Yan Da
2022-04-07 20:01:31 +08:00
parent 6b4da6f016
commit 040a2b6c75
5 changed files with 41 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
include/triton/ir/TritonOps.td
View File

@@ -123,8 +123,10 @@ def TT_LoadOp : TT_Op<"load", [SameOperandsAndResultShape]> {
let builders = [
// for args with default values
OpBuilder<(ins "Value":$ptr)>,
OpBuilder<(ins "Value":$ptr, "Value":$mask)>
OpBuilder<(ins "Value":$ptr, "triton::CacheModifier":$cache,
"triton::EvictionPolicy":$evict, "bool":$isVolatile)>,
OpBuilder<(ins "Value":$ptr, "Value":$mask, "triton::CacheModifier":$cache,
"triton::EvictionPolicy":$evict, "bool":$isVolatile)>
];
}
@@ -199,7 +201,7 @@ def TT_DotOp : TT_Op<"dot", [NoSideEffect, SameOperandsAndResultShape]> {
$d = matrix_multiply($a, $b) + $c
}];
let arguments = (ins TT_FpIntTensor:$a, TT_FpIntTensor:$b, TT_FpIntTensor:$c);
let arguments = (ins TT_FpIntTensor:$a, TT_FpIntTensor:$b, TT_FpIntTensor:$c, BoolAttr:$allowTF32);
let results = (outs TT_FpIntTensor:$d);
}

6
lib/ir/Ops.cpp
View File

@@ -33,7 +33,8 @@ void StoreOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, :
}
//-- LoadOp --
void LoadOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, ::mlir::Value ptr) {
void LoadOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, ::mlir::Value ptr,
::mlir::triton::CacheModifier cache, ::mlir::triton::EvictionPolicy evict, bool isVolatile) {
TensorType ptrType = ptr.getType().dyn_cast<TensorType>();
Type elementType = ptrType.getElementType().dyn_cast<PointerType>().getPointeeType();
auto shape = ptrType.getShape();
@@ -57,6 +58,9 @@ void LoadOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, ::
state.addOperands(ptr);
state.addOperands(mask);
state.addOperands(other);
state.addAttribute(cacheAttrName(state.name), ::mlir::triton::CacheModifierAttr::get(builder.getContext(), cache));
state.addAttribute(evictAttrName(state.name), ::mlir::triton::EvictionPolicyAttr::get(builder.getContext(), evict));
state.addAttribute(isVolatileAttrName(state.name), builder.getBoolAttr(isVolatile));
state.addTypes({resultType});
}

20
python/src/triton.cc
View File

@@ -1155,9 +1155,12 @@ void init_triton_ir(py::module &&m) {
return self.create<mlir::arith::OrIOp>(loc, lhs, rhs);
})
// // Input/Output
.def("create_load", [](mlir::OpBuilder &self, mlir::Value &ptrs) -> mlir::Value {
.def("create_load", [](mlir::OpBuilder &self, mlir::Value &ptrs,
mlir::triton::CacheModifier cacheModifer,
mlir::triton::EvictionPolicy evictionPolicy,
bool isVolatile) -> mlir::Value {
auto loc = self.getUnknownLoc();
return self.create<mlir::triton::LoadOp>(loc, ptrs);
return self.create<mlir::triton::LoadOp>(loc, ptrs, cacheModifer, evictionPolicy, isVolatile);
})
.def("create_store", [](mlir::OpBuilder &self, mlir::Value &ptrs, mlir::Value &value) -> void {
auto loc = self.getUnknownLoc();
@@ -1200,8 +1203,7 @@ void init_triton_ir(py::module &&m) {
})
.def("create_broadcast", [](mlir::OpBuilder &self, mlir::Value &arg, std::vector<int64_t> &shape) -> mlir::Value {
auto loc = self.getUnknownLoc();
// TODO: should be scalar type here
auto argType = arg.getType();
auto argType = arg.getType().dyn_cast<mlir::RankedTensorType>().getElementType();
return self.create<mlir::triton::BroadcastOp>(
loc, mlir::RankedTensorType::get(shape, argType), arg
);
@@ -1246,9 +1248,9 @@ void init_triton_ir(py::module &&m) {
loc, self.getI32Type(), self.getI32IntegerAttr(axis)
);
})
.def("create_dot", [](mlir::OpBuilder &self, mlir::Value &a, mlir::Value &b, mlir::Value &c) -> mlir::Value {
.def("create_dot", [](mlir::OpBuilder &self, mlir::Value &a, mlir::Value &b, mlir::Value &c, bool allowTF32) -> mlir::Value {
auto loc = self.getUnknownLoc();
return self.create<mlir::triton::DotOp>(loc, c.getType(), a, b, c);
return self.create<mlir::triton::DotOp>(loc, c.getType(), a, b, c, allowTF32);
})
// .def("create_exp", &ir::builder::create_exp, ret::reference)
// .def("create_cos", &ir::builder::create_cos, ret::reference)
@@ -1257,7 +1259,11 @@ void init_triton_ir(py::module &&m) {
// .def("create_trans", &ir::builder::create_trans, ret::reference)
// .def("create_sqrt", &ir::builder::create_sqrt, ret::reference)
// .def("create_reduce", &ir::builder::create_reduce, ret::reference)
// .def("create_select", &ir::builder::create_select, ret::reference)
.def("create_select", [](mlir::OpBuilder &self, mlir::Value &condition,
mlir::Value &trueValue, mlir::Value &falseValue) -> mlir::Value {
auto loc = self.getUnknownLoc();
return self.create<mlir::SelectOp>(loc, condition, trueValue, falseValue);
})
// // Intrinsics
// // These have no place in the IR, and hopefully they can be removed at some point
// .def("create_umulhi", &ir::builder::create_umulhi, ret::reference)

13
python/triton/code_gen.py
View File

@@ -1177,7 +1177,16 @@ class JITFunction:
# Compile to ttir, for the propose of testing MLIR rewriting
def compile_to_ttir(self, *wargs, grid, num_warps=4, num_stages=2, **kwargs):
# TODO: share code with _compile & __call__
# handle arguments passed by name
kwargs = {self.arg_names.index(name): value for name, value in kwargs.items()}
wargs = list(wargs)
for i, pos in enumerate(sorted(kwargs)):
wargs.insert(pos + i, kwargs[pos])
if len(wargs) != len(self.arg_names):
raise TypeError(f"Function takes {len(self.arg_names)} positional arguments but {len(wargs)} were given")
# handle annotations
for pos, _type in self.annotations.items():
wargs[pos] = _type(wargs[pos])
# preparing args
tensor_idxs = [i for i, arg in enumerate(wargs) if hasattr(arg, 'data_ptr')]
# attributes
@@ -1191,7 +1200,7 @@ class JITFunction:
attributes[i] = min(Kernel.pow2_divisor(addr),
Kernel.pow2_divisor(range_size))
# transforms ints whose value is one into constants for just-in-time compilation
constants = {i: arg for i, arg in enumerate(wargs) if isinstance(arg, int) and arg == 1 and i not in self.fn.do_not_specialize}
constants = {i: arg for i, arg in enumerate(wargs) if isinstance(arg, int) and arg == 1 and i not in self.do_not_specialize}
constants.update({i: arg.value for i, arg in enumerate(wargs) if isinstance(arg, triton.language.constexpr)})
constants.update({i: None for i, arg in enumerate(wargs) if arg is None})
arg_types = [Kernel._to_python_ir(arg) for i, arg in enumerate(wargs) if i not in constants]

14
rewrite-test/jit/vecadd.py
View File

@@ -1,3 +1,4 @@
from tarfile import BLOCKSIZE
import torch
import triton
import triton.language as tl
@@ -9,8 +10,8 @@ def add_kernel(
y_ptr, # *Pointer* to second input vector
output_ptr, # *Pointer* to output vector
n_elements, # Size of the vector
# BLOCK_SIZE: tl.constexpr, # Number of elements each program should process
# # NOTE: `constexpr` so it can be used as a shape value
BLOCK_SIZE: tl.constexpr, # Number of elements each program should process
# NOTE: `constexpr` so it can be used as a shape value
):
# There are multiple 'program's processing different data. We identify which program
# we are here
@@ -19,8 +20,8 @@ def add_kernel(
# for instance, if you had a vector of length 256 and block_size of 64, the programs
# would each access the elements [0:64, 64:128, 128:192, 192:256].
# Note that offsets is a list of pointers
block_start = pid * 256
offsets = block_start + tl.arange(0, 256)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
# Create a mask to guard memory operations against out-of-bounds accesses
mask = offsets < n_elements
# Load x and y from DRAM, masking out any extra elements in case the input is not a
@@ -37,7 +38,6 @@ y = torch.rand(size, device='cuda')
z = torch.empty_like(x)
# add_kernel[(1,)](x, y, z, size, 256)
# print(add_kernel[(1,)].kernel.compile_to_ttir())
mod, ctx = add_kernel.compile_to_ttir(x, y, z, size, grid=(1,))
mod.get_context()
# print(add_kernel.annotations)
mod, ctx = add_kernel.compile_to_ttir(x, y, z, size, BLOCK_SIZE=256, grid=(1,))
mod.dump()
# print(mod)