[mlir][bufferization] Support casts in EmptyTensorElimination

[springerm]

Summary

EmptyTensorElimination is a pre-bufferization transformation that replaces "tensor.empty" ops with "tensor.extract_slice" ops. This revision adds support for cases where the input IR contains "tensor.cast" ops.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct TraversalConfig {
   bool followInPlaceOnly = false;
 
   /// Specifies whether non-equivalent OpOperands should be followed.
   bool followEquivalentOnly = false;
 
   /// Specifies whether unknown/non-bufferizable/ops not included in the
   /// OpFilter of BufferizationOptions should be followed.
   bool followUnknownOps = false;
+
+  /// Specifies whether OpOperands with a different type that are not the result
+  /// of a CastOpInterface op should be followed.
+  bool followSameTypeOrCastsOnly = false;
 };
 
 /// AnalysisState provides a variety of helper functions for dealing with
 /// tensor values.
 class AnalysisState {
 public:
   /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return all tensor OpOperand* if the op is not

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

for (AliasingOpOperand a : aliases) {
       if (config.followInPlaceOnly && !isInPlace(*a.opOperand)) {
         // Stop iterating if `followInPlaceOnly` is set but the alias is
         // out-of-place.
         if (config.alwaysIncludeLeaves)
           result.insert(value);
         continue;
       }
 
+      if (config.followSameTypeOrCastsOnly &&
+          a.opOperand->get().getType() != value.getType() &&
+          !opResult.getDefiningOp<CastOpInterface>()) {
+        // Stop iterating if `followSameTypeOrCastsOnly` is set but the alias is
+        // has a different type and the op is not a cast.
+        if (config.alwaysIncludeLeaves)
+          result.insert(value);
+        continue;
+      }
+
       workingSet.insert(a.opOperand->get());
     }
   }
 
   return result;
 }
 
 // Find the values that define the contents of the given value.

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

for (OpOperand &operand : op->getOpOperands()) {
         continue;
 
       // Find tensor.empty ops on the reverse SSA use-def chain. Only follow
       // equivalent tensors. I.e., stop when there are ops such as extract_slice
       // on the path.
       TraversalConfig config;
       config.followEquivalentOnly = true;
       config.alwaysIncludeLeaves = false;
+      // Replace only if the types match or are static <-> dynamic casts. We do
+      // not support slices or reshapes.
+      // TODO: This could be extended to support IR such as:
+      // %0 = tensor.empty() : tensor<128xf32>
+      // %1 = "some_op"(%0) : (tensor<128xf32>) -> (tensor<128xf32>)
+      // %2 = tensor.expand_shape %1 ...
+      // %3 = tensor.insert_slice %2 into ...
+      config.followSameTypeOrCastsOnly = true;
       SetVector<Value> emptyTensors = state.findValueInReverseUseDefChain(
           operand.get(), /*condition=*/
           [&](Value val) { return val.getDefiningOp<tensor::EmptyOp>(); },
           config);
 
       for (Value v : emptyTensors) {
         Operation *emptyTensorOp = v.getDefiningOp();
 
-        // Replace only if the types match. We do not support slices or casts.
-        // TODO: This could be extended to support IR such as:
-        // %0 = tensor.empty() : tensor<128xf32>
-        // %1 = "some_op"(%0) : (tensor<128xf32>) -> (tensor<128xf32>)
-        // %2 = tensor.expand_shape %1 ...
-        // %3 = tensor.insert_slice %2 into ...
-        if (v.getType() != operand.get().getType())
-          continue;
-
         // Find a suitable insertion point. If no suitable insertion point for
         // the replacement can be found, skip this replacement.
         Operation *insertionPoint =
             findValidInsertionPoint(emptyTensorOp, neededValues);
         if (!insertionPoint)
           continue;
 
         rewriter.setInsertionPoint(insertionPoint);
         Value replacement =
             rewriteFunc(rewriter, emptyTensorOp->getLoc(), operand);
         if (!replacement)
           continue;
-
+        if (replacement.getType() != v.getType()) {
+          rewriter.setInsertionPointAfterValue(replacement);
+          replacement = rewriter.create<tensor::CastOp>(v.getLoc(), v.getType(),
+                                                        replacement);
+        }
         // Replace the tensor::EmptyOp.
         rewriter.replaceOp(emptyTensorOp, replacement);
         state.resetCache();
       }
     }
   });
 
   return success();

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir

func.func @insertion_point_outside_loop(%t : tensor<?xf32>, %sz : index,
   }
 
   return %r : tensor<?xf32>
 }
 
 // -----
 
 // EmptyTensorElimination does currently not apply to chains where the type is
-// changing. This test just ensures that we do not crash or generate IR that
-// does not verify.
+// changing. (Casts are supported.) This test just ensures that we do not crash
+// or generate IR that does not verify.
 
 // CHECK-LABEL: func @shape_mismatch
 func.func @shape_mismatch(%t: tensor<5x6x128xf32>) -> tensor<5x6x128xf32> {
   %cst = arith.constant 8.0 : f32
   %0 = tensor.empty() : tensor<128xf32>
   %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<128xf32>) -> tensor<128xf32>
   %2 = tensor.expand_shape %1 [[0, 1, 2]]
       : tensor<128xf32> into tensor<1x1x128xf32>
   %3 = tensor.insert_slice %2 into %t[2, 3, 0][1, 1, 128][1, 1, 1]
       : tensor<1x1x128xf32> into tensor<5x6x128xf32>
   return %3 : tensor<5x6x128xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @cast(
+//  CHECK-SAME:     %[[t:.*]]: memref<256xf32,
+//       CHECK:   %[[sv:.*]] = memref.subview %[[t]]
+//       CHECK:   linalg.fill {{.*}} outs(%[[sv]]
+//       CHECK:   return %[[t]]
+func.func @cast(%t: tensor<256xf32>) -> tensor<256xf32> {
+  %cst = arith.constant 8.0 : f32
+  %c128 = arith.constant 128 : index
+  %0 = tensor.empty(%c128) : tensor<?xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?xf32>) -> tensor<?xf32>
+  %2 = tensor.cast %1 : tensor<?xf32> to tensor<128xf32>
+  %3 = tensor.insert_slice %2 into %t[2][128][1]
+      : tensor<128xf32> into tensor<256xf32>
+  return %3 : tensor<256xf32>
+}
+
+// -----
+
 //      CHECK: func @parallel_insert_slice(
 // CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
 // CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
 func.func @parallel_insert_slice(
   %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
   %sz: index)
     -> (tensor<?xf32>)
 {