[FRONTEND] Fix issues with atomics (#849)

This commit is contained in:
Philippe Tillet
2022-11-06 20:52:11 -08:00
committed by GitHub
parent b6dbe959f0
commit a4ff0c362c
4 changed files with 32 additions and 5 deletions

View File

@@ -278,6 +278,20 @@ struct TritonStorePattern : public OpConversionPattern<triton::StoreOp> {
}
};
struct TritonAtomicRMWPattern
: public OpConversionPattern<triton::AtomicRMWOp> {
using OpConversionPattern<triton::AtomicRMWOp>::OpConversionPattern;
LogicalResult
matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
rewriter.replaceOpWithNewOp<triton::AtomicRMWOp>(
op, typeConverter->convertType(op.getType()), adaptor.atomic_rmw_op(),
adaptor.ptr(), adaptor.val(), adaptor.mask());
return success();
}
};
struct TritonExtElemwisePattern
: public OpConversionPattern<triton::ExtElemwiseOp> {
using OpConversionPattern<triton::ExtElemwiseOp>::OpConversionPattern;