[mlir][bufferize] Improve analysis of external functions

[springerm]

Summary

External functions have no body, so they cannot be analyzed. Assume conservatively that each tensor bbArg may be aliasing with each tensor result. Furthermore, assume that each function arg is read and written-to after bufferization. This default behavior can be controlled with bufferization.access (similar to bufferization.memory_layout) in test cases.

Also fix a bug in the dialect attribute verifier, which did not run for region argument attributes.

[springerm]

minor improvements

[springerm]

fix test

[aartbik]

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

196–218

for (int64_t idx = 0, e = funcOp.getFunctionType().getNumInputs(); idx < e; ++idx) {

[springerm]

address comments

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td

let description = [{
     found in [bufferization](/docs/Bufferization/) and [buffer
     deallocation](/docs/BufferDeallocationInternals/).
   }];
   let dependentDialects = [
     "AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
   ];
 
   let extraClassDeclaration = [{
+    /// Verify an attribute from this dialect on the argument at 'argIndex' for
+    /// the region at 'regionIndex' on the given operation. Returns failure if
+    /// the verification failed, success otherwise. This hook may optionally be
+    /// invoked from any operation containing a region.
+    LogicalResult verifyRegionArgAttribute(Operation *,
+                                           unsigned regionIndex,
+                                           unsigned argIndex,
+                                           NamedAttribute) override;
+
     /// An attribute that can override writability of buffers of tensor function
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kWritableAttrName = "bufferization.writable";
 
+    /// An attribute for function arguments that describes how the function
+    /// accesses the buffer. Can be one "none", "read", "write" or "read-write".
+    ///
+    /// When no attribute is specified, the analysis tries to infer the access
+    /// behavior from its body. In case of external functions, for which no
+    /// function body is available, "read-write" is assumed by default.
+    constexpr const static ::llvm::StringLiteral
+        kBufferAccessAttrName = "bufferization.access";
+
     /// Attribute name used to mark the bufferization layout for region
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kBufferLayoutAttrName = "bufferization.buffer_layout";
 
     /// Attribute name used to mark escaping behavior of buffer allocations.
     /// Escaping allocations cannot be deallocated in the same block and must
     /// be treated specially: They are currently deallocated with the

mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp

 void mlir::bufferization::BufferizationDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"
       >();
   addInterfaces<BufferizationInlinerInterface>();
 }
 
-LogicalResult
-BufferizationDialect::verifyOperationAttribute(Operation *op,
-                                               NamedAttribute attr) {
-  using bufferization::BufferizableOpInterface;
-
+LogicalResult BufferizationDialect::verifyRegionArgAttribute(
+    Operation *op, unsigned /*regionIndex*/, unsigned argIndex,
   if (attr.getName() == kWritableAttrName) {
     if (!attr.getValue().isa<BoolAttr>()) {
       return op->emitError() << "'" << kWritableAttrName
                              << "' is expected to be a boolean attribute";
     }
     if (!isa<FunctionOpInterface>(op))
-      return op->emitError() << "expected " << attr.getName()
-                             << " to be used on function-like operations";
+      return op->emitError() << "expected '" << kWritableAttrName
+                             << "' to be used on function-like operations";
+    if (cast<FunctionOpInterface>(op).isExternal())
+      return op->emitError() << "'" << kWritableAttrName
+                             << "' is invalid on external functions";
+    return success();
+  }
+  if (attr.getName() == kBufferAccessAttrName) {
+    if (!attr.getValue().isa<StringAttr>()) {
+      return op->emitError() << "'" << kBufferAccessAttrName
+                             << "' is expected to be a string attribute";
+    }
+    StringRef str = attr.getValue().cast<StringAttr>().getValue();
+    if (str != "none" && str != "read" && str != "write" && str != "read-write")
+      return op->emitError()
+             << "invalid value for '" << kBufferAccessAttrName << "'";
+    if (!isa<FunctionOpInterface>(op))
+      return op->emitError() << "expected '" << kBufferAccessAttrName
+                             << "' to be used on function-like operations";
     return success();
   }
   if (attr.getName() == kBufferLayoutAttrName) {
     if (!attr.getValue().isa<AffineMapAttr>()) {
       return op->emitError() << "'" << kBufferLayoutAttrName
                              << "' is expected to be a affine map attribute";
     }
     if (!isa<FunctionOpInterface>(op))
-      return op->emitError() << "expected " << attr.getName()
-                             << " to be used on function-like operations";
+      return op->emitError() << "expected '" << kBufferLayoutAttrName
+                             << "' to be used on function-like operations";
     return success();
   }
+  return op->emitError() << "attribute '" << kBufferLayoutAttrName
+                         << "' not supported as a region arg attribute by the "
+                            "bufferization dialect";
+}
+
+LogicalResult
+BufferizationDialect::verifyOperationAttribute(Operation *op,
+                                               NamedAttribute attr) {
+  using bufferization::BufferizableOpInterface;
+
   if (attr.getName() == kEscapeAttrName) {
     auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
     if (!arrayAttr)
       return op->emitError() << "'" << kEscapeAttrName
                              << "' is expected to be a bool array attribute";
     if (arrayAttr.size() != op->getNumResults())
       return op->emitError()
              << "'" << kEscapeAttrName
for (const auto &it : llvm::enumerate(arrayAttr)) {
                << "'" << kEscapeAttrName << "' only valid for tensor results";
       if (!bufferizableOp.bufferizesToAllocation(op->getOpResult(it.index())))
         return op->emitError() << "'" << kEscapeAttrName
                                << "' only valid for allocation results";
     }
     return success();
   }
 
-  return op->emitError() << "attribute '" << attr.getName()
-                         << "' not supported by the bufferization dialect";
+  return op->emitError()
+         << "attribute '" << attr.getName()
+         << "' not supported as an op attribute by the bufferization dialect";
 }

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

static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
 /// Store function BlockArguments that are equivalent to/aliasing a returned
 /// value in FuncAnalysisState.
 static LogicalResult
 aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                              FuncAnalysisState &funcState) {
+  if (funcOp.getBody().empty()) {
+    // No function body available. Conservatively assume that every tensor
+    // return value may alias with any tensor bbArg.
+    FunctionType type = funcOp.getFunctionType();
+    for (const auto &inputIt : llvm::enumerate(type.getInputs())) {
+      if (!inputIt.value().isa<TensorType>())
+        continue;
+      for (const auto &resultIt : llvm::enumerate(type.getResults())) {
+        if (!resultIt.value().isa<TensorType>())
+          continue;
+        int64_t returnIdx = resultIt.index();
+        int64_t bbArgIdx = inputIt.index();
+        funcState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
+        funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
+      }
+    }
+    return success();
+  }
+
   // Support only single return-terminated block in the function.
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<RankedTensorType>()) {
           int64_t returnIdx = returnVal.getOperandNumber();
           int64_t bbArgIdx = bbArg.getArgNumber();
           if (state.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
             funcState.equivalentFuncArgs[funcOp][returnIdx] = bbArgIdx;
             if (state.getOptions().testAnalysisOnly)
               annotateEquivalentReturnBbArg(returnVal, bbArg);
           }
           if (state.areAliasingBufferizedValues(returnVal.get(), bbArg)) {
             funcState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
             funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
           }
         }
 
   return success();
 }
 
-static void annotateFuncArgAccess(func::FuncOp funcOp, BlockArgument bbArg,
-                                  bool isRead, bool isWritten) {
+static void annotateFuncArgAccess(func::FuncOp funcOp, int64_t idx, bool isRead,
+                                  bool isWritten) {
   OpBuilder b(funcOp.getContext());
   Attribute accessType;
   if (isRead && isWritten) {
     accessType = b.getStringAttr("read-write");
   } else if (isRead) {
     accessType = b.getStringAttr("read");
   } else if (isWritten) {
     accessType = b.getStringAttr("write");
   } else {
     accessType = b.getStringAttr("none");
   }
-  funcOp.setArgAttr(bbArg.getArgNumber(), "bufferization.access", accessType);
+  funcOp.setArgAttr(idx, BufferizationDialect::kBufferAccessAttrName,
+                    accessType);
 }
 
 /// Determine which FuncOp bbArgs are read and which are written. When run on a
 /// function with unknown ops, we conservatively assume that such ops bufferize
 /// to a read + write.
 static LogicalResult
 funcOpBbArgReadWriteAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                              FuncAnalysisState &funcState) {
+  for (int64_t idx = 0, e = funcOp.getFunctionType().getNumInputs(); idx < e;
+       ++idx) {
+    // Skip non-tensor arguments.
+    if (!funcOp.getFunctionType().getInput(idx).isa<TensorType>())
+      continue;
+    bool isRead;
+    bool isWritten;
+    if (auto accessAttr = funcOp.getArgAttrOfType<StringAttr>(
+            idx, BufferizationDialect::kBufferAccessAttrName)) {
+      // Buffer access behavior is specified on the function. Skip the analysis.
+      StringRef str = accessAttr.getValue();
+      isRead = str == "read" || str == "read-write";
+      isWritten = str == "write" || str == "read-write";
+    } else if (funcOp.getBody().empty()) {
       // If the function has no body, conservatively assume that all args are
       // read + written.
-  if (funcOp.getBody().empty()) {
-    for (BlockArgument bbArg : funcOp.getArguments()) {
-      funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
-      funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
-    }
-
-    return success();
+      isRead = true;
+      isWritten = true;
+    } else {
+      // Analyze the body of the function.
+      BlockArgument bbArg = funcOp.getArgument(idx);
+      isRead = state.isValueRead(bbArg);
+      isWritten = state.isValueWritten(bbArg);
     }
 
-  for (BlockArgument bbArg : funcOp.getArguments()) {
-    if (!bbArg.getType().isa<TensorType>())
-      continue;
-    bool isRead = state.isValueRead(bbArg);
-    bool isWritten = state.isValueWritten(bbArg);
     if (state.getOptions().testAnalysisOnly)
-      annotateFuncArgAccess(funcOp, bbArg, isRead, isWritten);
+      annotateFuncArgAccess(funcOp, idx, isRead, isWritten);
     if (isRead)
-      funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.readBbArgs[funcOp].insert(idx);
     if (isWritten)
-      funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.writtenBbArgs[funcOp].insert(idx);
   }
 
   return success();
 }
 } // namespace
 
 /// Remove bufferization attributes on FuncOp arguments.
 static void removeBufferizationAttributes(BlockArgument bbArg) {
mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
-    // No body => no analysis.
-    if (funcOp.getBody().empty())
-      continue;
-
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, state, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state)))

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir

%r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
     %A2 = tensor.insert_slice %B2 into %t[%i][%sz][1] : tensor<?xf32> into tensor<?xf32>
     scf.yield %A2 : tensor<?xf32>
   }
   // CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true"]}
 
   return %r0 : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func private @ext_func(tensor<*xf32> {bufferization.access = "read-write"})
+func.func private @ext_func(%t: tensor<*xf32>)
+
+// CHECK: func.func @private_func_read_write(%{{.*}}: tensor<5xf32> {bufferization.access = "read"})
+func.func @private_func_read_write(%t: tensor<5xf32>) -> f32 {
+  %c0 = arith.constant 0 : index
+  // Bufferizes out-of-place because `ext_func` may modify the buffer.
+  // CHECK: tensor.cast {{.*}} {__inplace_operands_attr__ = ["false"]}
+  %0 = tensor.cast %t : tensor<5xf32> to tensor<*xf32>
+  func.call @ext_func(%0) : (tensor<*xf32>) -> ()
+  %1 = tensor.extract %t[%c0] : tensor<5xf32>
+  return %1 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @print_buffer(tensor<*xf32> {bufferization.access = "read"})
+func.func private @print_buffer(%t: tensor<*xf32> {bufferization.access = "read"})
+
+// CHECK: func.func @private_func_read(%{{.*}}: tensor<5xf32> {bufferization.access = "read"})
+func.func @private_func_read(%t: tensor<5xf32>) -> f32 {
+  %c0 = arith.constant 0 : index
+  // Bufferizes in-place because `print_buffer` is read-only.
+  // CHECK: tensor.cast {{.*}} {__inplace_operands_attr__ = ["true"]}
+  %0 = tensor.cast %t : tensor<5xf32> to tensor<*xf32>
+  // CHECK: call @print_buffer(%cast) {__inplace_operands_attr__ = ["true"]}
+  func.call @print_buffer(%0) : (tensor<*xf32>) -> ()
+  %1 = tensor.extract %t[%c0] : tensor<5xf32>
+  return %1 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @ext_func(tensor<?xf32> {bufferization.access = "read-write"}, tensor<?xf32> {bufferization.access = "read-write"})
+func.func private @ext_func(%t1: tensor<?xf32>, %t2: tensor<?xf32>)
+
+// CHECK: func.func @private_func_two_params_writing(%{{.*}}: tensor<?xf32> {bufferization.access = "read"})
+func.func @private_func_two_params_writing(%t: tensor<?xf32>) {
+  // Both operands bufferize out-of-place because both bufferize to a memory
+  // write.
+  // CHECK: call @ext_func(%{{.*}}, %{{.*}}) {__inplace_operands_attr__ = ["false", "false"]}
+  func.call @ext_func(%t, %t) : (tensor<?xf32>, tensor<?xf32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @ext_func(tensor<?xf32> {bufferization.access = "read-write"}) -> (tensor<5xf32>, tensor<6xf32>)
+func.func private @ext_func(%t: tensor<?xf32>) -> (tensor<5xf32>, tensor<6xf32>)
+
+// CHECK: func.func @private_func_aliasing(%{{.*}}: tensor<?xf32> {bufferization.access = "read"})
+func.func @private_func_aliasing(%t: tensor<?xf32>) -> f32 {
+  %c0 = arith.constant 0 : index
+  // Bufferizes out-of-place because either one of the two reuslts may alias
+  // with the argument and one of the results is read afterwards.
+  // CHECK: call @ext_func(%{{.*}}) {__inplace_operands_attr__ = ["false"]} : (tensor<?xf32>) -> (tensor<5xf32>, tensor<6xf32>)
+  %0, %1 = func.call @ext_func(%t) : (tensor<?xf32>) -> (tensor<5xf32>, tensor<6xf32>)
+  %2 = tensor.extract %1[%c0] : tensor<6xf32>
+  return %2 : f32
+}

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

func.func @scf_while_non_equiv_yield(%arg0: tensor<5xi1>,
 
   %v0 = tensor.extract %r0[%idx] : tensor<5xi1>
   %v1 = tensor.extract %r1[%idx] : tensor<5xi1>
   return %v0, %v1 : i1, i1
 }
 
 // -----
 
-func.func private @fun_with_side_effects(%A: tensor<?xf32> {bufferization.writable = true})
+func.func private @fun_with_side_effects(%A: tensor<?xf32>)
 
 func.func @foo(%A: tensor<?xf32> {bufferization.writable = true}) -> (tensor<?xf32>) {
   call @fun_with_side_effects(%A) : (tensor<?xf32>) -> ()
   return %A: tensor<?xf32>
 }
 
 func.func @scf_yield_needs_copy(%A : tensor<?xf32> {bufferization.writable = true}, %iters : index) {
   %c0 = arith.constant 0 : index

mlir/test/Dialect/Bufferization/invalid.mlir

 func.func private @foo(tensor<20x40xf32, #DCSR>) -> ()
 
 func.func @sparse_alloc_call() {
   // expected-error @+1{{sparse tensor allocation should not escape function}}
   %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
   call @foo(%0) : (tensor<20x40xf32, #DCSR>) -> ()
   return
 }
+
+// -----
+
+// expected-error @+1{{invalid value for 'bufferization.access'}}
+func.func private @invalid_buffer_access_type(tensor<*xf32> {bufferization.access = "foo"})
+
+// -----
+
+// expected-error @+1{{'bufferization.writable' is invalid on external functions}}
+func.func private @invalid_writable_attribute(tensor<*xf32> {bufferization.writable = false})
+
+// -----
+
+func.func @invalid_writable_on_op() {
+  // expected-error @+1{{attribute '"bufferization.writable"' not supported as an op attribute by the bufferization dialect}}
+  arith.constant {bufferization.writable = true} 0  : index
+}

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

func.func @scf_for_with_tensor.insert_slice(
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @execute_region_with_conflict(
 //  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
 func.func @execute_region_with_conflict(
-    %t1 : tensor<?xf32> {bufferization.writable = "true"})
+    %t1 : tensor<?xf32> {bufferization.writable = true})
   -> (f32, tensor<?xf32>, f32)
 {
   %f1 = arith.constant 0.0 : f32
   %idx = arith.constant 7 : index
 
   // scf.execute_region is canonicalized away after bufferization. So just the
   // memref.store is left over.