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

[aartbik]

LGTM, but please wait for Matthias to chime in

[springerm]

mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp

202–220

Makes sense to call these from SparsificationAndBufferizationPass but ideally these should be in OneShotModuleBufferize.cpp and OneShotModuleBufferize.h. These attributes are needed only for module bufferization.

[Peiming]

address comments.

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h

 };
 
 // Create a global op for the given tensor-valued constant in the program.
 // Globals are created lazily at the top of the enclosing ModuleOp with pretty
 // names. Duplicates are avoided.
 FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp constantOp,
                                          uint64_t alignment);
 
+// Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
+void removeBufferizationAttributesInModule(ModuleOp moduleOp);
+
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERUTILS_H

mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp

 //===- BufferUtils.cpp - buffer transformation utilities ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for buffer optimization passes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallString.h"
auto global = globalBuilder.create<memref::GlobalOp>(
       /*constant=*/true,
       /*alignment=*/memrefAlignment);
   symbolTable.insert(global);
   // The symbol table inserts at the end of the module, but globals are a bit
   // nicer if they are at the beginning.
   global->moveBefore(&moduleOp.front());
   return global;
 }
+
+//===----------------------------------------------------------------------===//
+// BufferizationAttributesCleanUp
+//===----------------------------------------------------------------------===//
+
+/// Remove bufferization attributes on FuncOp arguments.
+static void removeBufferizationAttributes(BlockArgument bbArg) {
+  auto funcOp = cast<func::FuncOp>(bbArg.getOwner()->getParentOp());
+  funcOp.removeArgAttr(bbArg.getArgNumber(),
+                       BufferizationDialect::kBufferLayoutAttrName);
+  funcOp.removeArgAttr(bbArg.getArgNumber(),
+                       BufferizationDialect::kWritableAttrName);
+}
+
+void bufferization::removeBufferizationAttributesInModule(ModuleOp moduleOp) {
+  moduleOp.walk([&](func::FuncOp op) {
+    for (BlockArgument bbArg : op.getArguments())
+      removeBufferizationAttributes(bbArg);
+  });
+}

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

 //
 // TODO: Add FuncOp attributes so that bbArgs of external FuncOps can be marked
 // as "not reading" and/or "not writing".
 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
 
for (BlockArgument bbArg : funcOp.getArguments()) {
     if (isWritten)
       funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
   }
 
   return success();
 }
 } // namespace
 
-/// Remove bufferization attributes on FuncOp arguments.
-static void removeBufferizationAttributes(BlockArgument bbArg) {
-  auto funcOp = cast<func::FuncOp>(bbArg.getOwner()->getParentOp());
-  funcOp.removeArgAttr(bbArg.getArgNumber(),
-                       BufferizationDialect::kBufferLayoutAttrName);
-  funcOp.removeArgAttr(bbArg.getArgNumber(),
-                       BufferizationDialect::kWritableAttrName);
-}
-
 /// Return the func::FuncOp called by `callOp`.
 static func::FuncOp getCalledFunction(CallOpInterface callOp) {
   SymbolRefAttr sym = callOp.getCallableForCallee().dyn_cast<SymbolRefAttr>();
   if (!sym)
     return nullptr;
   return dyn_cast_or_null<func::FuncOp>(
       SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
         LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
-  moduleOp.walk([&](func::FuncOp op) {
-    for (BlockArgument bbArg : op.getArguments())
-      removeBufferizationAttributes(bbArg);
-  });
+  removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }
 
 LogicalResult mlir::bufferization::runOneShotModuleBufferize(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

 //===- SparsificationAndBufferizationPass.cpp - Tensor to Memref Lowering -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 
denseOpFilter.allowOperation([&](Operation *op) {
       if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
         FunctionType funcType = funcOp.getFunctionType();
         if (containsSparseTensor(funcType.getInputs()) ||
             containsSparseTensor(funcType.getResults()))
           return false;
       }
       return true;
     });
-    return bufferization::bufferizeOp(getOperation(), bufferizationOptions,
+
+    if (failed(bufferization::bufferizeOp(getOperation(), bufferizationOptions,
                                           /*copyBeforeWrite=*/false,
-                                      &denseOpFilter);
+                                          &denseOpFilter)))
+      return failure();
+
+    bufferization::removeBufferizationAttributesInModule(getOperation());
+    return success();
   }
 
   void runOnOperation() override {
     {
       // Run enabling transformations.
       OpPassManager pm("builtin.module");
       pm.addPass(createPreSparsificationRewritePass());
       if (failed(runPipeline(pm, getOperation())))