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

Documentation

Browse Source
This commit is contained in:
Yan Da
2022-04-07 16:00:53 +08:00
parent 16d44e5c4c
commit 6b4da6f016
3 changed files with 26 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
include/triton/ir/Traits.h
View File

@@ -17,13 +17,14 @@ class TensorSizeTrait : public TraitBase<ConcreteType, TensorSizeTrait> {
public:
// TODO: move impl to .cc files
static LogicalResult verifyTrait(Operation *op) {
int constexpr maxElement = 1048576;
for (auto opType : op->getOperandTypes()) {
if (auto tensorType = opType.dyn_cast<RankedTensorType>()) {
int64_t numElements = 1;
for (int64_t s : tensorType.getShape())
numElements *= s;
if (numElements > 1048576)
return op->emitError("Maximum allowed number of elements is 1048576, but ")
if (numElements > maxElement)
return op->emitError("Maximum allowed number of elements is ") << maxElement << ", but "
<< *op << " has more than that";
if ((numElements & (numElements - 1)) != 0)
return op->emitError("Number of elements must be power-of-two, but ")
@@ -36,8 +37,8 @@ public:
int64_t numElements = 1;
for (int64_t s : tensorType.getShape())
numElements *= s;
if (numElements > 1048576)
return op->emitError("Maximum allowed number of elements is 1048576, but ")
if (numElements > maxElement)
return op->emitError("Maximum allowed number of elements is ") << maxElement << ", but "
<< *op << " has more than that";
if ((numElements & (numElements - 1)) != 0)
return op->emitError("Number of elements must be power-of-two, but ")

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

@@ -195,6 +195,10 @@ def TT_GetNumProgramsOp : TT_Op<"get_num_programs"> {
def TT_DotOp : TT_Op<"dot", [NoSideEffect, SameOperandsAndResultShape]> {
let summary = "dot";
let description = [{
$d = matrix_multiply($a, $b) + $c
}];
let arguments = (ins TT_FpIntTensor:$a, TT_FpIntTensor:$b, TT_FpIntTensor:$c);
let results = (outs TT_FpIntTensor:$d);
@@ -279,6 +283,12 @@ def TT_AtomicCASOp : TT_Op<"atomic_cas"> {
def TT_MakeRangeOp : TT_Op<"make_range", [NoSideEffect]> {
let summary = "make range";
let description = [{
Returns an 1D int32 tensor.
Values span from $start to $end (exclusive), with step = 1
}];
let arguments = (ins I32Attr:$start, I32Attr:$end);
let results = (outs TT_IntegerTensor:$result);