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

More progress on TritonGPU conversion

Browse Source
This commit is contained in:
Yan Da
2022-05-04 14:54:31 +08:00
parent 3ad7bee35e
commit b9279d2e3b
4 changed files with 48 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
lib/Conversion/TritonToTritonGPU/TritonToTritonGPU.cpp
View File

@@ -26,15 +26,15 @@ public:
}
};
template<class Op>
class ArithCmpPattern : public OpConversionPattern<Op> {
template<class SrcOp, class DstOp>
class ArithCmpPattern : public OpConversionPattern<SrcOp> {
public:
using OpConversionPattern<Op>::OpConversionPattern;
using OpConversionPattern<SrcOp>::OpConversionPattern;
LogicalResult matchAndRewrite(Op op, typename Op::Adaptor adaptor,
LogicalResult matchAndRewrite(SrcOp op, typename SrcOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Type retType = this->getTypeConverter()->convertType(op.getType());
Op res = rewriter.replaceOpWithNewOp<Op>(
DstOp res = rewriter.replaceOpWithNewOp<DstOp>(
op, retType, adaptor.getPredicate(), adaptor.getLhs(), adaptor.getRhs()
);
return success();
@@ -106,8 +106,10 @@ void populateArithmeticPatternsAndLegality(
ArithBinaryPattern<arith::DivFOp>,
ArithBinaryPattern<arith::RemFOp>,
// Cmp
ArithCmpPattern<arith::CmpIOp>,
ArithCmpPattern<arith::CmpFOp>
// ArithCmpPattern<arith::CmpIOp, triton::gpu::CmpIOp>,
// ArithCmpPattern<arith::CmpFOp, triton::gpu::CmpFOp>
ArithCmpPattern<arith::CmpIOp, triton::gpu::CmpIOp>,
ArithCmpPattern<arith::CmpFOp, triton::gpu::CmpFOp>
>(typeConverter, context);
}

6
lib/Dialect/Triton/IR/Ops.cpp
View File

@@ -13,14 +13,14 @@ namespace triton {
static Type getI1SameShape(Type type) {
auto i1Type = IntegerType::get(type.getContext(), 1);
if (auto tensorType = type.dyn_cast<RankedTensorType>())
return RankedTensorType::get(tensorType.getShape(), i1Type);
return RankedTensorType::get(tensorType.getShape(), i1Type, tensorType.getEncoding());
return Type();
}
static Type getI32SameShape(Type type) {
auto i32Type = IntegerType::get(type.getContext(), 32);
if (auto tensorType = type.dyn_cast<RankedTensorType>())
return RankedTensorType::get(tensorType.getShape(), i32Type);
return RankedTensorType::get(tensorType.getShape(), i32Type, tensorType.getEncoding());
return Type();
}
@@ -29,7 +29,7 @@ static Type getPointerTypeFromTensor(Type type) {
Type elementType = tensorType.getElementType();
auto shape = tensorType.getShape();
PointerType ptrType = PointerType::get(elementType, 1);
return RankedTensorType::get(shape, ptrType);
return RankedTensorType::get(shape, ptrType, tensorType.getEncoding());
}
return Type();
}

24
lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp
View File

@@ -53,35 +53,27 @@ TritonGPUTypeConverter::TritonGPUTypeConverter(MLIRContext *context,
TritonGPUConversionTarget::TritonGPUConversionTarget(
MLIRContext &context, TritonGPUTypeConverter &typeConverter)
: ConversionTarget(context), typeConverter(typeConverter) {
addLegalDialect<triton::TritonDialect,
StandardOpsDialect,
scf::SCFDialect>();
addLegalDialect<StandardOpsDialect, scf::SCFDialect>();
// Some ops from SCF are illegal
addIllegalOp<scf::ExecuteRegionOp, scf::ParallelOp,
scf::ReduceOp, scf::ReduceReturnOp>();
addDynamicallyLegalDialect<arith::ArithmeticDialect>([&](Operation *op) {
if (typeConverter.isLegal(op)) {
// llvm::errs() << *op << " is dyanamically legal\n";
if (typeConverter.isLegal(op))
return true;
}
// if (typeConverter.isLegal(op->getOperandTypes()))
// llvm::errs() << *op << " is illegal with legal operands\n";
// if (typeConverter.isLegal(op->getResultTypes())) {
// llvm::errs() << *op << " is illegal with legal results\n";
// llvm::errs() << "operand0: " << op->getOperand(0) << "\n"
// << "operand1: " << op->getOperand(1) << "\n";
// }
return false;
});
addDynamicallyLegalDialect<triton::TritonDialect>([&](Operation *op) {
if (typeConverter.isLegal(op))
return true;
// llvm::errs() << *op << " is illegal\n"
// << "inside ...\n"
// << *op->getParentOp() << "\n";
return false;
});
addDynamicallyLegalDialect<triton::gpu::TritonGPUDialect>([&](Operation *op) {
if (typeConverter.isLegal(op))
return true;
return false;
});