[Mlir-commits] [mlir] 3230a64 - [mlir][NVGPU] Fix incorrect API usage in RewritePatterns

[Matthias Springer]

Author: Matthias Springer
Date: 2023-03-02T16:15:09+01:00
New Revision: 3230a64936d28d1114e9493e620ebbd65f595e07

URL: https://github.com/llvm/llvm-project/commit/3230a64936d28d1114e9493e620ebbd65f595e07
DIFF: https://github.com/llvm/llvm-project/commit/3230a64936d28d1114e9493e620ebbd65f595e07.diff

LOG: [mlir][NVGPU] Fix incorrect API usage in RewritePatterns

Incorrect API usage was detected by D144552.

Differential Revision: https://reviews.llvm.org/D145156

Added: 
    

Modified: 
    mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp b/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp
index 7430daf9e5fec..292738de4b52a 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp
@@ -38,7 +38,7 @@ struct MmaSyncF32ToTF32Pattern : public OpRewritePattern<nvgpu::MmaSyncOp> {
         precision(precision) {}
 
   LogicalResult matchAndRewrite(nvgpu::MmaSyncOp op,
-                                PatternRewriter &rewrite) const override {
+                                PatternRewriter &rewriter) const override {
     Location location = op->getLoc();
 
     if (op->hasAttr(op.getTf32EnabledAttrName()) ||
@@ -53,8 +53,10 @@ struct MmaSyncF32ToTF32Pattern : public OpRewritePattern<nvgpu::MmaSyncOp> {
       return emitError(location, "TF32x3 is not supported at the moment "
                                  "for nvgpu.mma.sync on f32 datatype");
 
-    if (precision == MmaSyncF32Lowering::TF32)
-      op.setTf32EnabledAttr(rewrite.getUnitAttr());
+    if (precision == MmaSyncF32Lowering::TF32) {
+      rewriter.updateRootInPlace(
+          op, [&]() { op.setTf32EnabledAttr(rewriter.getUnitAttr()); });
+    }
 
     return success();
   }