[mlir][bufferize] Do not insert useless casts for newly allocated buffers

[springerm]

Summary

Newly allocated buffers do not have a layout map. Bufferization patterns can add a fully dynamic layout map via casting. This should be done only if the layout map is actually needed and not by default on every allocation.

When bufferizing scf.for ops, such layout maps may actually be needed. E.g., after bufferizing an scf.for loop:

scf.for ... -> memref<?xf32, #map> {
  // We alloc because of an out-of-place bufferization.
  %0 = memref.alloc : memref<?xf32>
  // Here we actually need a cast.
  scf.yield %0 : memref<?xf32>
}

Depends On D123321

[akuegel]

Also just curious, does this really depend on the other patch?

[springerm]

In D123369#3438441, @akuegel wrote:

Also just curious, does this really depend on the other patch?

Yes, it does. The diffbase commit fixes an issue with missing arith.select casts that it triggered when removing the unneeded casts in here.

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

 FailureOr<Value> BufferizationState::createAlloc(OpBuilder &b, Location loc,
                                                  Value shapedValue,
                                                  Optional<bool> dealloc) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // Compute allocation memref type.
   assert(shapedValue.getType().isa<ShapedType>());
-  MemRefType memRefType = shapedValue.getType().dyn_cast<MemRefType>();
   SmallVector<Value> dynShape;
   MemRefType allocMemRefType =
       getAllocationTypeAndShape(b, loc, shapedValue, dynShape);
 
   // Should be the buffer be deallocated again or should we let it leak?
   bool skipDealloc;
   if (dealloc) {
     skipDealloc = !dealloc.getValue();
   } else {
     assert(shapedValue.getType().isa<TensorType>() &&
            "must specify `dealloc` if non-tensor value is passed");
     // Buffer should be not be deallocated if deallocs are generally deactivated
     // or if the tensor is yielded from a block.
     skipDealloc = !getOptions().createDeallocs ||
                   getAnalysisState().isTensorYielded(shapedValue);
   }
 
   // Create the buffer allocation.
-  Value alloc =
-      createBufferAllocation(b, loc, allocMemRefType, dynShape, skipDealloc);
-
-  // Insert a cast if a different type was requested.
-  if (memRefType && memRefType != allocMemRefType) {
-    assert(memref::CastOp::areCastCompatible(allocMemRefType, memRefType) &&
-           "createAlloc: cast incompatible");
-    alloc = b.create<memref::CastOp>(loc, memRefType, alloc);
-  }
-
-  return alloc;
+  return createBufferAllocation(b, loc, allocMemRefType, dynShape, skipDealloc);
 }
 
 /// Create a memory copy between two memref buffers.
 LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
                                           Value from, Value to,
                                           const BufferizationOptions &options) {
   if (options.memCpyFn)
     return (*options.memCpyFn)(b, loc, from, to);

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

SmallVector<Value> yieldValues =
           // TODO: We should rollback, but for now just assume that this always
           // succeeds.
           assert(yieldedAlloc.hasValue() && "could not create alloc");
           LogicalResult copyStatus =
               bufferization::createMemCpy(rewriter, val.getLoc(), yieldedVal,
                                           *yieldedAlloc, state.getOptions());
           (void)copyStatus;
           assert(succeeded(copyStatus) && "could not create memcpy");
+
+          if (yieldedVal.getType() == yieldedAlloc->getType())
             return *yieldedAlloc;
+
+          // The iter_arg memref type has a layout map. Cast the new buffer to
+          // the same type.
+          // TODO: In case the iter_arg has a layout map that is not the fully
+          // dynamic one, we cannot cast the new buffer. In that case, the
+          // iter_arg must be changed to the fully dynamic layout map. (And then
+          // the new buffer can be casted.)
+          assert(memref::CastOp::areCastCompatible(yieldedAlloc->getType(),
+                                                   yieldedVal.getType()) &&
+                 "scf.for op bufferization: cast incompatible");
+          Value casted = rewriter.create<memref::CastOp>(
+              val.getLoc(), yieldedVal.getType(), *yieldedAlloc);
+          return casted;
         });
     yieldOp.getResultsMutable().assign(yieldValues);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 
     return success();
   }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir

 // -----
 
 // CHECK-TENSOR-LABEL: func @simple_tensor_test(
 //  CHECK-TENSOR-SAME:     %[[t1:.*]]: tensor<?xf32>
 func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
   // CHECK-TENSOR: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
   %c0 = arith.constant 0 : index
   // CHECK-TENSOR: %[[alloc:.*]] = memref.alloc
-  // CHECK-TENSOR: %[[casted:.*]] = memref.cast %[[alloc]]
-  // CHECK-TENSOR: %[[casted_tensor:.*]] = bufferization.to_tensor %[[casted]]
+  // CHECK-TENSOR: %[[casted_alloc:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK-TENSOR: memref.copy %[[t1_memref]], %[[alloc]]
   // CHECK-TENSOR: memref.store %{{.*}}, %[[alloc]]
   %0 = tensor.insert %f into %t1[%c0] : tensor<?xf32>
-  // CHECK-TENSOR: return %[[casted_tensor]]
+  // CHECK-TENSOR: return %[[casted_alloc]]
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-SCF-LABEL: func @simple_scf_if(
 //  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32> {linalg.inplaceable = true}, %[[c:.*]]: i1, %[[pos:.*]]: index
 func @simple_scf_if(%t1: tensor<?xf32> {linalg.inplaceable = true}, %c: i1, %pos: index, %f: f32)

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

 // CHECK-LABEL: func @return_tensor(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
 func @return_tensor(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
   %c0 = arith.constant 0 : index
 
   // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
   // CHECK: %[[dim:.*]] = tensor.dim %[[A]]
   // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
-  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: memref.copy %[[A_memref]], %[[alloc]]
   // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
-  // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[casted]]
+  // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
   %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
 
   // CHECK: memref.dealloc %[[alloc]]
   // CHECK: return %[[res_tensor]]
   return %0 : tensor<?xf32>
 }
 
 // -----

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

 //===----------------------------------------------------------------------===//
 // Simple loop cases
 //===----------------------------------------------------------------------===//
 
 // -----
 
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
-// CHECK-LABEL: func @scf_for_yield_only
-//  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
+// CHECK-LABEL: func @scf_for_yield_only(
+//  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>,
 //  CHECK-SAME:   %[[t:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
+//  CHECK-SAME:   ) -> memref<?xf32> {
 func @scf_for_yield_only(%A : tensor<?xf32> {linalg.inplaceable = false},
                          %B : tensor<?xf32> {linalg.inplaceable = true},
                          %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   //     CHECK:   %[[ALLOC_FOR_A:.*]] = memref.alloc
-  //     CHECK:   %[[CASTED:.*]] = memref.cast %[[ALLOC_FOR_A]]
   //     CHECK:   memref.copy %[[A]], %[[ALLOC_FOR_A]]
 
   // The first scf.for remains but just turns into dead code.
   %r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     scf.yield %t : tensor<?xf32>
   }
 
   // The second scf.for remains but just turns into dead code.
   %r1 = scf.for %i = %lb to %ub step %step iter_args(%t = %B) -> (tensor<?xf32>) {
     scf.yield %t : tensor<?xf32>
   }
 
-  //     CHECK:   return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
+  //     CHECK:   return %[[ALLOC_FOR_A]] : memref<?xf32>
   // CHECK-NOT:   dealloc
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // Ensure that the function bufferizes without error. This tests pre-order
 // traversal of scf.for loops during bufferization. No need to check the IR,

 func @scf_for_with_tensor.insert_slice(
    %A : tensor<?xf32> {linalg.inplaceable = false},
    %B : tensor<?xf32> {linalg.inplaceable = true},
    %C : tensor<4xf32> {linalg.inplaceable = false},
    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   //     CHECK:   %[[ALLOC_FOR_A:.*]] = memref.alloc
-  //     CHECK:   %[[CASTED:.*]] = memref.cast %[[ALLOC_FOR_A]]
   //     CHECK:   memref.copy %[[A]], %[[ALLOC_FOR_A]]
 
   //     CHECK: %[[svA:.*]] = memref.subview %[[ALLOC_FOR_A]][0] [4] [1]
   //     CHECK: %[[svB:.*]] = memref.subview %[[B]][0] [4] [1]
 
   //     CHECK:   scf.for {{.*}}
   // CHECK-NOT: iter_args
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     // %ttA bufferizes to direct copy of %BUFFER_CAST_C into %svA
     //     CHECK: memref.copy %[[C]], %[[svA]]
     %ttA = tensor.insert_slice %C into %tA[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
     // %ttB bufferizes to direct copy of %BUFFER_CAST_C into %BUFFER_CAST_B
     //     CHECK:   memref.copy %[[C]], %[[svB]]
     %ttB = tensor.insert_slice %C into %tB[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
     // CHECK-NOT:   scf.yield
     scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
   }
 
-  //     CHECK:  return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
+  //     CHECK:  return %[[ALLOC_FOR_A]] : memref<?xf32>
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Cross function boundary cases.
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @execute_region_with_conflict(
 //  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
 func @execute_region_with_conflict(%t1 : tensor<?xf32> {linalg.inplaceable = "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.
 
   // CHECK: %[[alloc:.*]] = memref.alloc
-  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: memref.copy %[[m1]], %[[alloc]]
   // CHECK: memref.store %{{.*}}, %[[alloc]][%{{.*}}]
   %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) {
     %t2 = tensor.insert %f1 into %t1[%idx] : tensor<?xf32>
     scf.yield %f1, %t2, %f1 : f32, tensor<?xf32>, f32
   }
 
+  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: %[[load:.*]] = memref.load %[[m1]]
   %3 = tensor.extract %t1[%idx] : tensor<?xf32>
 
   // CHECK: return %{{.*}}, %[[casted]], %[[load]] : f32, memref<?xf32, #{{.*}}>, f32
   return %0, %1, %3 : f32, tensor<?xf32>, f32
 }
 
 // -----
func @write_after_select_read_one(
     %t2 : tensor<?xf32> {linalg.inplaceable = true},
     %c : i1)
   -> (f32, tensor<?xf32>)
 {
   %cst = arith.constant 0.0 : f32
   %idx = arith.constant 0 : index
 
   // CHECK: %[[alloc:.*]] = memref.alloc
-  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
-  // CHECK: memref.copy %[[t1]], %[[alloc]]
+  // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK-DAG: memref.copy %[[t1]], %[[alloc]]
   // CHECK: %[[select:.*]] = arith.select %{{.*}}, %[[casted]], %[[t2]]
   %s = arith.select %c, %t1, %t2 : tensor<?xf32>
 
   // CHECK: memref.store %{{.*}}, %[[select]]
   %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
 
   // CHECK: %[[f:.*]] = memref.load %[[t1]]
   %f = tensor.extract %t1[%idx] : tensor<?xf32>

 
 // Note: This bufferizes to inefficient code, but bufferization should not see
 // such IR in the first place. The iter_arg would canonicalize away. This test
 // case is just to ensure that the bufferization generates correct code.
 
 // CHECK-LABEL: func @scf_for_yield_non_equivalent(
 //  CHECK-SAME:     %[[t:.*]]: memref<?xf32
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}})
-//       CHECK:   %[[casted:.*]] = memref.cast %[[alloc]]
-//       CHECK:   %[[for:.*]] = scf.for {{.*}} iter_args(%[[iter:.*]] = %[[casted]])
+//       CHECK:   %[[for:.*]] = scf.for {{.*}} iter_args(%[[iter:.*]] = %[[alloc]])
 //       CHECK:     memref.dealloc %[[iter]]
 //       CHECK:     %[[alloc2:.*]] = memref.alloc(%{{.*}})
 //       CHECK:     memref.copy %[[t]], %[[alloc2]]
-//       CHECK:     %[[casted2:.*]] = memref.cast %[[alloc2]]
-//       CHECK:     scf.yield %[[casted2]]
+//       CHECK:     scf.yield %[[alloc2]]
 //       CHECK:   return %[[for]]
 func @scf_for_yield_non_equivalent(
     %t: tensor<?xf32>, %lb : index, %ub : index, %step : index) -> tensor<?xf32> {
   %r = scf.for %i = %lb to %ub step %step iter_args(%a = %t) -> tensor<?xf32> {
     scf.yield %t : tensor<?xf32>
   }
 
   return %r : tensor<?xf32>

mlir/test/Dialect/Linalg/one-shot-module-bufferize.mlir

 //  CHECK-DAG:   call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]]
   %r0:2 = call @scf_for_with_tensor_insert_slice(%A, %B, %C, %lb, %ub, %step) :
       (tensor<?xf32>, tensor<?xf32>, tensor<4xf32>, index, index, index)
         -> (tensor<?xf32>, tensor<?xf32>)
 
   // %r0#0 requires a copy because we have no idea what the function is doing.
 //  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
 //  CHECK-DAG:   %[[casted:.*]] = memref.cast %[[alloc]]
-//      CHECK:   memref.copy %[[B]], %[[alloc]]
+//  CHECK-DAG:   memref.copy %[[B]], %[[alloc]]
 // CHECK-NEXT:   call @some_external_func(%[[casted]]) : (memref<?xf32, #[[$DYN_1D_MAP]]>) -> ()
   call @some_external_func(%r0#0) : (tensor<?xf32>) -> ()
 
 //      CHECK:   return
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----

 func @entry(%A : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
             %B : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
             %C : tensor<?xf32> {linalg.inplaceable = false}) {
 // Note: `callee` does not write to its bbArg directly, but `external_func`
 // does. Inside `callee`, the writes via `external_func` do not cause a
 // conflict. However, inside `entry`, the writes do cause a conflict because
 // %A, %B and %C are not inplaceable. This test case shows that this kind of
 // conflict detection has a "transitive" nature.
-//      CHECK: %[[ALLOC_C:.*]] = memref.alloc
-//      CHECK: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]]
-//      CHECK: %[[ALLOC_B:.*]] = memref.alloc
-//      CHECK: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]]
-//      CHECK: %[[ALLOC_A:.*]] = memref.alloc
-//      CHECK: memref.copy %[[A]], %[[ALLOC_A]]
-//      CHECK: memref.copy %[[B]], %[[ALLOC_B]]
-//      CHECK: memref.copy %[[C]], %[[ALLOC_C]]
-//      CHECK: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]]
+//  CHECK-DAG: %[[ALLOC_C:.*]] = memref.alloc
+//  CHECK-DAG: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]]
+//  CHECK-DAG: %[[ALLOC_B:.*]] = memref.alloc
+//  CHECK-DAG: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]]
+//  CHECK-DAG: %[[ALLOC_A:.*]] = memref.alloc
+//  CHECK-DAG: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]]
+//  CHECK-DAG: memref.copy %[[A]], %[[ALLOC_A]]
+//  CHECK-DAG: memref.copy %[[B]], %[[ALLOC_B]]
+//  CHECK-DAG: memref.copy %[[C]], %[[ALLOC_C]]
 // CHECK-NEXT: call @callee(%[[CASTED_A]], %[[CASTED_B]], %[[CASTED_C]])
   call @callee(%A, %B, %C) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> ()
   return
 }
 
 // -----
 
 // No alloc or copy inside of the loop.

 
 // CHECK-LABEL: func @equivalent_func_arg_2(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
 func @equivalent_func_arg_2(%t0: tensor<?xf32> {linalg.inplaceable = true},
                             %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
   // CHECK: scf.for {{.*}} {
   %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
     // CHECK: %[[alloc:.*]] = memref.alloc
-    // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
-    // CHECK: memref.copy %[[arg0]], %[[alloc]]
+    // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
+    // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]
     // CHECK: call @inner_func_2(%[[casted]])
     // CHECK: memref.dealloc %[[alloc]]
     // CHECK-NOT: scf.yield
     %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
     scf.yield %t1 : tensor<?xf32>
   }
   return %1: tensor<?xf32>
 }