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

Fix ReduceOp conversion

Browse Source
This commit is contained in:
Yan Da
2022-05-25 16:03:06 +08:00
parent 9b670cfb9f
commit e6f89a5777
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -224,6 +224,19 @@ struct TritonGenericPattern : public OpConversionPattern<Op> {
} }
}; };
struct TritonReducePattern : public OpConversionPattern<triton::ReduceOp> {
using OpConversionPattern<triton::ReduceOp>::OpConversionPattern;
LogicalResult matchAndRewrite(triton::ReduceOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Type retType = this->getTypeConverter()->convertType(op.getType());
auto newOp = rewriter.replaceOpWithNewOp<triton::ReduceOp>(
op, retType, adaptor.redOp(), adaptor.operand(), adaptor.axis()
);
return success();
}
};
void populateTritonPatterns( void populateTritonPatterns(
TritonGPUTypeConverter& typeConverter, RewritePatternSet &patterns TritonGPUTypeConverter& typeConverter, RewritePatternSet &patterns
) { ) {
@@ -231,6 +244,7 @@ void populateTritonPatterns(
patterns.add<TritonGenericPattern<triton::ReshapeOp>, patterns.add<TritonGenericPattern<triton::ReshapeOp>,
TritonGenericPattern<triton::BroadcastOp>, TritonGenericPattern<triton::BroadcastOp>,
TritonGenericPattern<triton::GEPOp>, TritonGenericPattern<triton::GEPOp>,
TritonReducePattern,
TritonMakeRangePattern, TritonMakeRangePattern,
TritonDotPattern, TritonDotPattern,
TritonLoadPattern, TritonLoadPattern,