[Mlir-commits] [mlir] c7a9e5e - [mlir][sparse][bufferization] cleanup bufferization attributes after SparsificationAndBufferizationPass

[Peiming Liu]

Author: Peiming Liu
Date: 2022-12-02T23:03:54Z
New Revision: c7a9e5e5d28ded8cd0f7162956d0df679c341392

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

LOG: [mlir][sparse][bufferization] cleanup bufferization attributes after SparsificationAndBufferizationPass

Reviewed By: aartbik, springerm

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
    mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
index aafa002f42061..f6402b0726c2f 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
@@ -30,6 +30,9 @@ LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state);
 LogicalResult bufferizeModuleOp(ModuleOp moduleOp,
                                 const OneShotBufferizationOptions &options);
 
+/// Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
+void removeBufferizationAttributesInModule(ModuleOp moduleOp);
+
 /// Run One-Shot Module Bufferization on the given module. Performs a simple
 /// function call analysis to determine which function arguments are
 /// inplaceable. Then analyzes and bufferizes FuncOps one-by-one with One-Shot

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 8b70c2af76ab5..e2878b2e1df76 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -377,6 +377,14 @@ mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
   return success();
 }
 
+void mlir::bufferization::removeBufferizationAttributesInModule(
+    ModuleOp moduleOp) {
+  moduleOp.walk([&](func::FuncOp op) {
+    for (BlockArgument bbArg : op.getArguments())
+      removeBufferizationAttributes(bbArg);
+  });
+}
+
 LogicalResult mlir::bufferization::bufferizeModuleOp(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options) {
   assert(options.bufferizeFunctionBoundaries &&
@@ -405,10 +413,7 @@ LogicalResult mlir::bufferization::bufferizeModuleOp(
   }
 
   // Post-pass cleanup of function argument attributes.
-  moduleOp.walk([&](func::FuncOp op) {
-    for (BlockArgument bbArg : op.getArguments())
-      removeBufferizationAttributes(bbArg);
-  });
+  removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index b3d10813407c6..6066a17e27914 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
@@ -75,9 +76,14 @@ class SparsificationAndBufferizationPass
       }
       return true;
     });
-    return bufferization::bufferizeOp(getOperation(), bufferizationOptions,
-                                      /*copyBeforeWrite=*/false,
-                                      &denseOpFilter);
+
+    if (failed(bufferization::bufferizeOp(getOperation(), bufferizationOptions,
+                                          /*copyBeforeWrite=*/false,
+                                          &denseOpFilter)))
+      return failure();
+
+    bufferization::removeBufferizationAttributesInModule(getOperation());
+    return success();
   }
 
   void runOnOperation() override {