20100a7254
This PR merges the `triton-mlir` branch, in which we have been quietly rewriting the Triton backend from scratch to increase maintainability, stability and ultimately performance. Changes to the runtime are minimal, and this new version aims to remain backward-compatible with the previous commit. The legacy backend is now officially deprecated, but can still be accessed via the `legacy-backend` tag. Co-authored-by: Keren Zhou <kerenzhou@openai.com> Co-authored-by: Yan Chunwei <yanchunwei@outlook.com> Co-authored-by: goostavz <109190422+goostavz@users.noreply.github.com> Co-authored-by: Shintaro Iwasaki <siwasaki@fb.com> Co-authored-by: Yan Da <dyanab@connect.ust.hk> Co-authored-by: Jun Yang <yangjunpro@gmail.com> Co-authored-by: Ian Bearman <ianb@microsoft.com> Co-authored-by: Jason Ansel <jansel@jansel.net> Co-authored-by: Qingyi Liu <qingyil@nvidia.com> Co-authored-by: ben-zhang-609 <110140741+ben-zhang-609@users.noreply.github.com> Co-authored-by: Chenggang Zhao <lyricz@yeah.net> Co-authored-by: ben-zhang-609 <benzh609@gmail.com> Co-authored-by: dongdongl <dongdongl@nvidia.com>
72 lines
2.7 KiB
C++
72 lines
2.7 KiB
C++
#include "triton/Dialect/Triton/IR/Traits.h"
|
|
|
|
static mlir::LogicalResult verifySameEncoding(mlir::Type tyA, mlir::Type tyB) {
|
|
using namespace mlir;
|
|
auto encA = tyA.dyn_cast<RankedTensorType>();
|
|
auto encB = tyA.dyn_cast<RankedTensorType>();
|
|
if (!encA || !encB)
|
|
return success();
|
|
return encA.getEncoding() == encB.getEncoding() ? success() : failure();
|
|
}
|
|
|
|
mlir::LogicalResult
|
|
mlir::OpTrait::impl::verifySameOperandsAndResultEncoding(Operation *op) {
|
|
if (failed(verifyAtLeastNOperands(op, 1)) ||
|
|
failed(verifyAtLeastNResults(op, 1)))
|
|
return failure();
|
|
|
|
auto type = op->getOperand(0).getType();
|
|
for (auto resultType : op->getResultTypes())
|
|
if (failed(verifySameEncoding(resultType, type)))
|
|
return op->emitOpError()
|
|
<< "requires the same encoding for all operands and results";
|
|
return verifySameOperandsEncoding(op);
|
|
}
|
|
|
|
mlir::LogicalResult
|
|
mlir::OpTrait::impl::verifySameOperandsEncoding(Operation *op) {
|
|
if (failed(verifyAtLeastNOperands(op, 1)))
|
|
return failure();
|
|
|
|
auto type = op->getOperand(0).getType();
|
|
for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1))
|
|
if (failed(verifySameEncoding(opType, type)))
|
|
return op->emitOpError() << "requires the same encoding for all operands";
|
|
|
|
return success();
|
|
}
|
|
|
|
mlir::LogicalResult mlir::OpTrait::impl::verifyTensorSize(Operation *op) {
|
|
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 > maxTensorNumElements)
|
|
return op->emitError("Maximum allowed number of elements is ")
|
|
<< maxTensorNumElements << ", but " << *op
|
|
<< " has more than that";
|
|
if ((numElements & (numElements - 1)) != 0)
|
|
return op->emitError("Number of elements must be power-of-two, but ")
|
|
<< *op << " doesn't follow the rule (" << numElements << ")"
|
|
<< " elements";
|
|
}
|
|
}
|
|
for (auto opType : op->getResultTypes()) {
|
|
if (auto tensorType = opType.dyn_cast<RankedTensorType>()) {
|
|
int64_t numElements = 1;
|
|
for (int64_t s : tensorType.getShape())
|
|
numElements *= s;
|
|
if (numElements > maxTensorNumElements)
|
|
return op->emitError("Maximum allowed number of elements is ")
|
|
<< maxTensorNumElements << ", but " << *op
|
|
<< " has more than that";
|
|
if ((numElements & (numElements - 1)) != 0)
|
|
return op->emitError("Number of elements must be power-of-two, but ")
|
|
<< *op << " doesn't follow the rule (" << numElements << ")"
|
|
<< " elements";
|
|
}
|
|
}
|
|
return success();
|
|
}
|