[mlir][linalg][bufferize][NFC] Print results of FuncOp read/write analysis

[springerm]

Summary

Print more information with test-analysis-only.

Depends On D118891

[springerm]

rebase

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

#endif // NDEBUG
   });
   return isWritten;
 }
 
 /// Determine which FuncOp bbArgs are read and which are written. If this
 /// PostAnalysisStep is run on a function with unknown ops, it will
 /// conservatively assume that such ops bufferize to a read + write.
 struct FuncOpBbArgReadWriteAnalysis : public PostAnalysisStep {
+  static void annotateFuncArg(FuncOp funcOp, BlockArgument bbArg, bool isRead,
+                              bool isWritten) {
+    OpBuilder b(funcOp.getContext());
+    Attribute accessType;
+    if (isRead && isWritten) {
+      accessType = b.getStringAttr("read-write");
+    } else if (isRead) {
+      accessType = b.getStringAttr("read");
+    } else if (isWritten) {
+      accessType = b.getStringAttr("write");
+    } else {
+      accessType = b.getStringAttr("none");
+    }
+    funcOp.setArgAttr(bbArg.getArgNumber(), "bufferization.access", accessType);
+  }
+
   LogicalResult run(Operation *op, BufferizationState &state,
                     BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override {
     ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
     auto funcOp = cast<FuncOp>(op);
 
     // If the function has no body, conservatively assume that all args are
     // read + written.
     if (funcOp.getBody().empty()) {
       for (BlockArgument bbArg : funcOp.getArguments()) {
         moduleState.readBbArgs.insert(bbArg);
         moduleState.writtenBbArgs.insert(bbArg);
       }
 
       return success();
     }
 
     for (BlockArgument bbArg : funcOp.getArguments()) {
       if (!bbArg.getType().isa<TensorType>())
         continue;
-      if (state.isValueRead(bbArg))
+      bool isRead = state.isValueRead(bbArg);
+      bool isWritten = isValueWritten(bbArg, state, aliasInfo);
+      if (state.getOptions().testAnalysisOnly)
+        annotateFuncArg(funcOp, bbArg, isRead, isWritten);
+      if (isRead)
         moduleState.readBbArgs.insert(bbArg);
-      if (isValueWritten(bbArg, state, aliasInfo))
+      if (isWritten)
         moduleState.writtenBbArgs.insert(bbArg);
     }
 
     return success();
   }
 };
 } // namespace
 

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

 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="test-analysis-only allow-return-memref" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="test-analysis-only allow-return-memref analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="test-analysis-only allow-return-memref analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="test-analysis-only allow-return-memref analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 //===----------------------------------------------------------------------===//
 // Simple cases
 //===----------------------------------------------------------------------===//
 
 // -----
 
-// CHECK-LABEL: func @extract_slice_fun
+// CHECK-LABEL: func @extract_slice_fun(
 func @extract_slice_fun(%A : tensor<?xf32> {linalg.inplaceable = false},
+//  CHECK-SAME:          bufferization.access = "read"
                         %B : tensor<?xf32> {linalg.inplaceable = true})
+//  CHECK-SAME:         bufferization.access = "read"
   -> (tensor<4xf32>, tensor<8xf32>)
 {
   // tensor.extract_slice is not used in a write, it is not compelled to
   // bufferize out of place. Let callers decide whether they want to create
   // aliasing subviews at all call sites or whether they allocate.
   // This is true irrespective of whether the function argument is inplaceable.
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
   %r0 = tensor.extract_slice %A[0][4][1] : tensor<?xf32> to tensor<4xf32>
 
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
   %r1 = tensor.extract_slice %B[0][8][1] : tensor<?xf32> to tensor<8xf32>
 
   return %r0, %r1: tensor<4xf32>, tensor<8xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @insert_slice_fun
+// CHECK-LABEL: func @insert_slice_fun(
 func @insert_slice_fun(%A : tensor<?xf32> {linalg.inplaceable = false},
+//  CHECK-SAME:        bufferization.access = "read"
                        %B : tensor<?xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:        bufferization.access = "read-write"
                        %C : tensor<4xf32> {linalg.inplaceable = false})
+//  CHECK-SAME:        bufferization.access = "read"
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   // must bufferize out of place.
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_operands_attr__ = ["true", "false"]}
   %r0 = tensor.insert_slice %C into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   // bufferizes inplace.
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
   %r1 = tensor.insert_slice %C into %B[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   //      CHECK: return
   // CHECK-SAME: __equivalent_func_args__ = [-1, 1]
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @conflict_on_B
+// CHECK-LABEL: func @conflict_on_B(
 func @conflict_on_B(%A : tensor<4x4xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:     bufferization.access = "read"
                     %B : tensor<4x4xf32> {linalg.inplaceable = true})
+//  CHECK-SAME:     bufferization.access = "read-write"
   -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
 {
   // matmul output operand interferes with input operand.
   //     CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
   %C = linalg.matmul  ins(%A, %B: tensor<4x4xf32>, tensor<4x4xf32>)
                      outs(%B: tensor<4x4xf32>)
     -> tensor<4x4xf32>

 }
 
 //===----------------------------------------------------------------------===//
 // Length-1 producer-consumer cases.
 //===----------------------------------------------------------------------===//
 
 // -----
 
-// CHECK-LABEL: func @extract_slice_extract_slice
+// CHECK-LABEL: func @extract_slice_extract_slice(
 func @extract_slice_extract_slice(
     %A : tensor<?xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:         bufferization.access = "read"
     %B : tensor<?xf32> {linalg.inplaceable = false})
+//  CHECK-SAME:         bufferization.access = "read"
   -> (tensor<2xf32>, tensor<2xf32>)
 {
   // tensor.extract_slice is not used in a write, it is not compelled to
   // bufferize out of place. Let callers decide whether they want to create
   // aliasing subviews at all call sites or whether they allocate.
   // This is true irrespective of whether the function argument is inplaceable.
   // CHECK: {__inplace_operands_attr__ = ["true"]}
   %r0 = tensor.extract_slice %A[0][4][1] : tensor<?xf32> to tensor<4xf32>
 
   // CHECK: {__inplace_operands_attr__ = ["true"]}
   %r1 = tensor.extract_slice %r0[0][2][1] : tensor<4xf32> to tensor<2xf32>
 
   // CHECK: {__inplace_operands_attr__ = ["true"]}
   %r2 = tensor.extract_slice %B[0][4][1] : tensor<?xf32> to tensor<4xf32>
 
   // CHECK: {__inplace_operands_attr__ = ["true"]}
   %r3 = tensor.extract_slice %r2[0][2][1] : tensor<4xf32> to tensor<2xf32>
 
   return %r1, %r3: tensor<2xf32>, tensor<2xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @insert_slice_insert_slice
+// CHECK-LABEL: func @insert_slice_insert_slice(
 func @insert_slice_insert_slice(
     %A : tensor<?xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:         bufferization.access = "read-write"
     %A2 : tensor<4xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:          bufferization.access = "read-write"
     %A3 : tensor<2xf32> {linalg.inplaceable = true},
+//  CHECK-SAME:          bufferization.access = "read"
     %B : tensor<?xf32> {linalg.inplaceable = false},
+//  CHECK-SAME:         bufferization.access = "read"
     %B2 : tensor<4xf32> {linalg.inplaceable = false},
+//  CHECK-SAME:          bufferization.access = "read"
     %B3 : tensor<2xf32> {linalg.inplaceable = false})
+//  CHECK-SAME:          bufferization.access = "read"
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   // CHECK: {__inplace_operands_attr__ = ["true", "true"]}
   %r0 = tensor.insert_slice %A3 into %A2[0][2][1] : tensor<2xf32> into tensor<4xf32>
 
   // CHECK: {__inplace_operands_attr__ = ["true", "true"]}
   %r1 = tensor.insert_slice %r0 into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 

 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Chain of tensor.insert_slice is better traversed in reverse order without
 // prioritizing  the tensor.insert_slice ops.
 //===----------------------------------------------------------------------===//
 
+// CHECK-LABEL: func @insert_slice_chain(
 func @insert_slice_chain(
     %v1: vector<32x90xf32>,
     %v2: vector<30x90xf32>,
     %arg0: tensor<62x126xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+// CHECK-SAME: bufferization.access = "none"
     %arg1: tensor<126x90xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+// CHECK-SAME: bufferization.access = "none"
     %arg2: tensor<62x90xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = true})
+// CHECK-SAME: bufferization.access = "write"
   -> tensor<62x90xf32> attributes {passthrough = [["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]}
 {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32
 
   //      CHECK: linalg.fill
   // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]
   %0 = linalg.fill(%cst, %arg2) : f32, tensor<62x90xf32> -> tensor<62x90xf32>
#accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>
 ]
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel"]
 }
 
-// CHECK-LABEL: func @linalg_op_same_out_tensors
+// CHECK-LABEL: func @linalg_op_same_out_tensors(
 func @linalg_op_same_out_tensors(
     %t1: tensor<?xf32> {linalg.inplaceable = true},
-    %t2: tensor<?xf32> {linalg.inplaceable = true}) -> (tensor<?xf32>, tensor<?xf32>){
+// CHECK-SAME:          bufferization.access = "read-write"
+    %t2: tensor<?xf32> {linalg.inplaceable = true})
+// CHECK-SAME:          bufferization.access = "write"
+  -> (tensor<?xf32>, tensor<?xf32>){
 
   //      CHECK: linalg.generic
   // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]
   %o:2 = linalg.generic #trait ins(%t1 : tensor<?xf32>)
                                outs (%t2, %t2 : tensor<?xf32>, tensor<?xf32>) {
       ^bb(%0: f32, %1: f32, %2 : f32) :
         linalg.yield %0, %0 : f32, f32
     } -> (tensor<?xf32>, tensor<?xf32>)
#accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>
 ]
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel"]
 }
 
-// CHECK-LABEL: func @linalg_op_same_out_tensors_2
+// CHECK-LABEL: func @linalg_op_same_out_tensors_2(
 func @linalg_op_same_out_tensors_2(
     %t1: tensor<?xf32> {linalg.inplaceable = true},
+// CHECK-SAME:          bufferization.access = "read-write"
     %t2: tensor<?xf32> {linalg.inplaceable = true})
+// CHECK-SAME:          bufferization.access = "write"
         -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>){
 
   //      CHECK: linalg.generic
   // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true", "false"]
   %o:3 = linalg.generic #trait
           ins(%t1 : tensor<?xf32>)
           outs (%t2, %t2, %t2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {
       ^bb(%0: f32, %1: f32, %2 : f32, %3 : f32) :