[Mlir-commits] [mlir] f74f091 - [mlir][linalg][bufferize] Relax tensor.insert_slice conflict rules

[Matthias Springer]

Author: Matthias Springer
Date: 2021-11-10T18:23:29+09:00
New Revision: f74f09128bc2fafe1d39a95a0c066731f854cbc1

URL: https://github.com/llvm/llvm-project/commit/f74f09128bc2fafe1d39a95a0c066731f854cbc1
DIFF: https://github.com/llvm/llvm-project/commit/f74f09128bc2fafe1d39a95a0c066731f854cbc1.diff

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

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.

Differential Revision: https://reviews.llvm.org/D113446

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index cf9ad6b5418a..8dda1b2f40b3 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -903,6 +903,31 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
           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;

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
index f2c56c3f75f2..fc82522c3ef5 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
@@ -213,6 +213,72 @@ func @extract_slice_matching_insert_slice(
 
 // -----
 
+// 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>,
@@ -946,6 +1012,28 @@ func @interleaved_extract_insert_slice_chain_2(
 
 // -----
 
+// 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)>
 ]