[mlir][linalg][bufferize] Relax tensor.insert_slice conflict rules

[springerm]

Summary

A tensor.insert_slice write does not conflict with a subsequent read of the source if the source is originating from a matching tensor.extract_slice.

[mravishankar]

LGTM.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

924

Why not use areEquivalentExtractSliceOps?

[springerm]

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

924

areEquivalentExtractSliceOps is called internally by hasMatchingExtractSliceOp.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

for (OpOperand *uConflictingWrite : usesWrite) {
           // InsertSliceOp is writing.
           //
           // In the above example:
           // uRead             = OpOperand 0 (%1) of tensor.insert_slice
           // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
           continue;
       }
 
+      // If uConflictingWrite is an InsertSliceOp...
+      if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp))
+        // As an example, consider the following IR.
+        //
+        // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace= [true] }
+        // %1 = linalg.fill %cst, %0 {inplace= [true] }
+        // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
+        //     {inplace= [true] }
+        // %3 = vector.transfer_read %1, %cst
+        //
+        // In the above example:
+        // uRead             = OpOperand 0 (%1) of vector.transfer_read
+        // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
+        // lastWrite         = %1
+        //
+        // This is not a conflict because the InsertSliceOp overwrites the
+        // memory segment of %1 with the exact same data. (Effectively, there
+        // is no memory write here.)
+        if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
+            aliasInfo.areEquivalentBufferizedValues(uRead->get(),
+                                                    insertSliceOp.source()) &&
+            hasMatchingExtractSliceOp(aliasInfo, insertSliceOp.source(),
+                                      insertSliceOp))
+          continue;
+
       // All requirements are met. Conflict found!
       LDBG("CONFLICT CONFIRMED!\n\n");
       return true;
     }
   }
 
   LDBG("NOT A CONFLICT!\n\n");
   return false;

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

func @extract_slice_matching_insert_slice(
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %r3 = tensor.insert_slice %r2 into %B[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @read_of_matching_insert_slice_source
+func @read_of_matching_insert_slice_source(
+    %A : tensor<?xf32> {linalg.inplaceable = true}, %idx : index, %idx2 : index)
+  -> (tensor<?xf32>, vector<5xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %cst2 = arith.constant 1.0 : f32
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32>
+
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
+
+  %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
+  return %2, %3 : tensor<?xf32>, vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @read_of_matching_insert_slice_source_interleaved
+func @read_of_matching_insert_slice_source_interleaved(
+    %A : tensor<?xf32> {linalg.inplaceable = true}, %idx : index, %idx2 : index,
+    %idx3 : index)
+  -> (tensor<?xf32>, vector<5xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %cst2 = arith.constant 1.0 : f32
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32>
+
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %4 = tensor.extract_slice %2[%idx3][%idx3][1] : tensor<?xf32> to tensor<?xf32>
+
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %5 = linalg.fill(%cst, %4) : f32, tensor<?xf32> -> tensor<?xf32>
+
+  %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %6 = tensor.insert_slice %5 into %2[%idx3][%idx3][1] : tensor<?xf32> into tensor<?xf32>
+
+  return %6, %3 : tensor<?xf32>, vector<5xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_slice_linalg_readonly_use
 func @extract_slice_linalg_readonly_use(
     %A : tensor<?x?xf32>,
     %B : tensor<4x4xf32>,
     %C : tensor<4x4xf32> {linalg.inplaceable = true})
   ->  (tensor<4x4xf32>, tensor<4x4xf32>)
 {
   // tensor.extract_slice is only used as a read, no interference irrespective
func @interleaved_extract_insert_slice_chain_2(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %10 into %8[31, 0] [30, 90] [1, 1] : tensor<30x90xf32> into tensor<62x90xf32>
 
   return %15 : tensor<62x90xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @extract_once_insert_twice
+func @extract_once_insert_twice(
+    %arg2: tensor<62x90xf32> {linalg.inplaceable = true})
+  -> (tensor<62x90xf32>)
+{
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+  %2 = tensor.extract_slice %arg2[0, 0] [32, 90] [1, 1] : tensor<62x90xf32> to tensor<32x90xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+  %8 = tensor.insert_slice %2 into %arg2[0, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+  %15 = tensor.insert_slice %2 into %8[15, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
+
+  return %15 : tensor<62x90xf32>
+}
+
+// -----
+
 #accesses = [
   affine_map<(i) -> (i)>
 ]
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel"]
 }