[mlir][bufferization] Add restrict and writable attrs to to_tensor

[springerm]

Summary

restrict is similar to the C++ restrict keyword. Such tensors are guaranteed to not alias with any tensors in the program (after bufferization).

[nicolasvasilache]

Nice!

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

def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
     SameOperandsAndResultShape,
     SameOperandsAndResultElementType,
     TypesMatchWith<"result type matches tensor equivalent of 'memref'",
                    "memref", "result",
                    "memref::getTensorTypeFromMemRefType($_self)">
   ]> {
   let summary = "memref to tensor operation";
   let description = [{
-    Create a tensor from a `memref`, making an independent copy of the element
-    data. The result value is a tensor whose shape and element type match the
-    memref operand.
+    An operation that creates a tensor from a `memref`. The result value is a
+    tensor whose shape and element type match the memref operand.
 
     The opposite of this op is `to_memref`. Together, these two ops are
     useful for source/target materializations when doing type conversions
     involving tensors and memrefs.
 
     Example:
 
     ```mlir
     // Produces a value of tensor<4x?xf32> type.
-    %12 = bufferization.to_tensor %10 : memref<4x?xf32, #layout, memspace0>
+    %t = bufferization.to_tensor %m : memref<4x?xf32, #layout, 0>
     ```
 
-    If tensor load is used in the bufferization steps, mutating the source
-    buffer after loading leads to undefined behavior.
+    If the `writable` unit attribute is set, the produced tensor is considered
+    "writable" during bufferization. Otherwise, every OpOperand that bufferizes
+    to a write to the future buffer of the resulting tensor (or an alias
+    thereof) will bufferize out-of-place to prevent emitting any writes to
+    `memref` during bufferization.
+
+    If the given memref does not alias with any other memref passed to another
+    `to_tensor` op, the `restrict` unit attribute can be set. Only such
+    operations are supported by One-Shot Bufferize. (Otherwise, potential memref
+    aliasing relationships would have to be captured in One-Shot Bufferize.)
+
+    Example:
+
+    ```
+    %t = bufferization.to_tensor %m restrict writable : memref<4xf32>
+
+    // %t is writable, so the tensor.insert may bufferize in-place in the
+    // absence of other conflicts.
+    %r = tensor.insert %f into %t[%idx] : tensor<4xf32>
+    ```
+
+    `to_tensor` ops are not bufferized. They are expected to fold away after
+    bufferization. If there are non-bufferizable ops in the IR and
+    `allowUnknownOps` is set, they may be part of the resulting IR and not fold
+    away. However, such IR is no longer bufferizable with One-Shot Bufferize.
   }];
 
   let arguments = (ins Arg<AnyRankedOrUnrankedMemRef,
-                       "the reference to load from", [MemRead]>:$memref);
+                       "the reference to load from", [MemRead]>:$memref,
+                       UnitAttr:$restrict, UnitAttr:$writable);
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     /// The result of a to_tensor is always a tensor.
     TensorType getType() {
       Type resultType = getResult().getType();
       if (resultType.isa<TensorType>())
         return resultType.cast<TensorType>();
       return {};
     }
 
     //===------------------------------------------------------------------===//
     // BufferizableOpInterface implementation
     //===------------------------------------------------------------------===//
 
-    // ToTensorOp conceptually loads a tensor from a memory location. The
-    // One-Shot analysis has no information about the memref that is loaded from
-    // by ToTensorOp. We have to assume that the loaded tensor may after
-    // bufferization potentially alias with any other bufferized tensor. Since
-    // ToTensorOp and ToMemrefOp have no aliasing OpOperand/OpResult pairs, this
-    // cannot be encoded directly in the analysis. However, declaring ToTensorOp
-    // results as not writable enforces a buffer copy and has the same effect.
-
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options) const {
-      // to_tensor cannot be bufferized. However, other ops that are using
-      // to_tensor's result will eventually be bufferized. At that point, they
-      // will start using to_tensor's memref operand. Once all users of
-      // to_tensor are bufferized, the op will not have any users anymore and
-      // DCE away. In case of partial bufferization, to_memref(to_tensor(x))
-      // constructs may be left over. These are folded by the canonicalizer or
-      // FinalizingBufferize.
+      // to_tensor/to_memref pairs fold away after bufferization.
       return success();
     }
 
-    bool isWritable(Value value, const AnalysisState &state) const {
-      // It is unknown whether the memref operand is writable or not.
-      return false;
-    }
+    bool isWritable(Value value, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
         Value value, const BufferizationOptions &options,
         const DenseMap<Value, BaseMemRefType> &fixedTypes) {
       return getMemref().getType().cast<BaseMemRefType>();
     }
   }];
 
-  let assemblyFormat = "$memref attr-dict `:` type($memref)";
+  let assemblyFormat = [{
+    $memref (`restrict` $restrict^)? (`writable` $writable^)? attr-dict
+      `:` type($memref)
+  }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 
 //===----------------------------------------------------------------------===//
 // ToMemrefOp
 //===----------------------------------------------------------------------===//
 
 def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
     BufferizableOpInterface,
     SameOperandsAndResultShape,
     SameOperandsAndResultElementType,
     Pure,
     TypesMatchWith<"type of 'tensor' is the tensor equivalent of 'memref'",
                    "memref", "tensor",
                    "memref::getTensorTypeFromMemRefType($_self)">
   ]> {
   let summary = "tensor to memref cast operation";
   let description = [{
-    Casts a tensor to a memref.
+    An operation that returns the future buffer of a `tensor`.
 
     ```mlir
-    // Result type is memref<4x?xf32, #layout, 42>
-    %12 = bufferization.to_memref %10 : memref<4x?xf32, #layout, 42>
+    // Result type is memref<4x?xf32, #layout, 0>
+    %m = bufferization.to_memref %t : memref<4x?xf32, #layout, 0>
     ```
 
-    Note, that mutating the result of the `to_memref` operation leads to
-    undefined behavior.
-
     This operation is a specialized variant of the built-in
-    `unrealized_conversion_cast` and is intended for use in the context of
-    gradual bufferization.
+    `unrealized_conversion_cast` and is used to make sure that the IR stays
+    valid at any point during the bufferization.
+
+    IR that contains `to_memref` ops cannot be bufferized with One-Shot
+    Bufferize.
   }];
 
   let arguments = (ins AnyTensor:$tensor);
   let results = (outs AnyRankedOrUnrankedMemRef:$memref);
 
   let extraClassDeclaration = [{
     //===------------------------------------------------------------------===//
     // BufferizableOpInterface implementation

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

LogicalResult DeallocTensorOp::bufferize(RewriterBase &rewriter,
   rewriter.eraseOp(getOperation());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ToTensorOp
 //===----------------------------------------------------------------------===//
 
+bool ToTensorOp::isWritable(Value value, const AnalysisState &state) {
+  return getWritable();
+}
+
 OpFoldResult ToTensorOp::fold(FoldAdaptor) {
   if (auto toMemref = getMemref().getDefiningOp<ToMemrefOp>())
     // Approximate alias analysis by conservatively folding only when no there
     // is no interleaved operation.
     if (toMemref->getBlock() == this->getOperation()->getBlock() &&
         toMemref->getNextNode() == this->getOperation())
       return toMemref.getTensor();
   return {};

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

static LogicalResult checkAliasInfoConsistency(Operation *op,
   WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
     if (!options.isOpAllowed(op.getOperation()))
       return WalkResult::advance();
 
     // Input IR may not contain any ToMemrefOps. These are not supported because
     // the analysis cannot follow the data flow through memrefs.
     if (isa<ToMemrefOp>(op.getOperation())) {
-      op->emitError("to_memref ops not supported during One-Shot Analysis");
+      op->emitError("to_memref ops are not supported by One-Shot Analysis");
       return WalkResult::interrupt();
     }
 
+    // Input IR may not contain any ToTensorOps without the "restrict"
+    // attribute. Such tensors may alias any other tensor, which is currently
+    // not handled in the analysis.
+    if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {
+      if (!toTensorOp.getRestrict()) {
+        op->emitError("to_tensor ops without `restrict` are not supported by "
+                      "One-Shot Analysis");
+        return WalkResult::interrupt();
+      }
+    }
+
     for (OpOperand &opOperand : op->getOpOperands()) {
       if (opOperand.get().getType().isa<TensorType>()) {
         if (wouldCreateReadAfterWriteInterference(
                 opOperand, domInfo, state,
                 /*checkConsistencyOnly=*/true)) {
           // This error can happen if certain "mustBufferizeInPlace" interface
           // methods are implemented incorrectly, such that the IR already has
           // a RaW conflict before making any bufferization decisions.

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

func.func @main_func(%A : tensor<?xf32> {bufferization.writable = true},
   %0 = call @some_use(%A, %v) : (tensor<?xf32>, vector<5xf32>) -> (tensor<?xf32>)
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @to_tensor_op_not_writable
 func.func @to_tensor_op_not_writable(%m: memref<?xf32>, %v:  vector<5xf32>,
                                      %idx1: index, %idx2: index)
     -> vector<10xf32> {
-  %0 = bufferization.to_tensor %m : memref<?xf32>
+  %0 = bufferization.to_tensor %m restrict : memref<?xf32>
 
   // Write to the tensor. Cannot be inplace due to tensor_load.
   //      CHECK: vector.transfer_write
   // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]
   %w = vector.transfer_write %v, %0[%idx1] : vector<5xf32>, tensor<?xf32>
 
   // Read from the tensor and return result.
   %cst = arith.constant 0.0 : f32

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

func.func @main() -> tensor<4xi32> {
   }
 
   // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r: tensor<4xi32>
 }
 
 // -----
 
-func.func @to_memref_op_is_writing(
+func.func @to_memref_op_unsupported(
     %t1: tensor<?xf32> {bufferization.writable = true}, %idx1: index,
     %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) {
-  // This is a RaW conflict because to_memref is an inplace write and %t1 is
-  // read further down. This will likely have to change with partial
-  // bufferization.
 
-  // expected-error @+1 {{to_memref ops not supported during One-Shot Analysis}}
+  // expected-error @+1 {{to_memref ops are not supported by One-Shot Analysis}}
   %0 = bufferization.to_memref %t1 : memref<?xf32>
 
   // Read from both.
   %cst = arith.constant 0.0 : f32
   %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
   %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
 
   return %r1, %r2 : vector<5xf32>, vector<5xf32>
 }
 
 // -----
 
+func.func @to_tensor_op_unsupported(%m: memref<?xf32>, %idx: index) -> (f32) {
+  // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}}
+  %0 = bufferization.to_tensor %m : memref<?xf32>
+
+  %1 = tensor.extract %0[%idx] : tensor<?xf32>
+  return %1 : f32
+}
+
+// -----
+
 // expected-error @+2 {{failed to bufferize op}}
 // expected-error @+1 {{cannot bufferize bodiless function that returns a tensor}}
 func.func private @foo(%t : tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
 
 func.func @call_to_unknown_tensor_returning_func(%t : tensor<?xf32>) {
   call @foo(%t) : (tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
   return
 }

mlir/test/Dialect/Bufferization/ops.mlir

%0 = bufferization.to_memref %arg0
     : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>
   %1 = bufferization.to_memref %arg1
     : memref<*xi64, 1>
   return %0, %1 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>
 }
 
 // CHECK-LABEL: func @test_to_tensor
 func.func @test_to_tensor(%buf : memref<2xf32>) -> tensor<2xf32> {
-  %tensor = bufferization.to_tensor %buf : memref<2xf32>
+  %tensor = bufferization.to_tensor %buf restrict writable : memref<2xf32>
   return %tensor : tensor<2xf32>
 }
 
 // CHECK-LABEL: func @test_alloc_tensor_op
 func.func @test_alloc_tensor_op(%t: tensor<?x5xf32>, %sz: index)
   -> tensor<?x5xf32>
 {
   // CHECK: bufferization.alloc_tensor(%{{.*}}) : tensor<?x5xf32>