[Mlir-commits] [mlir] [ml_program] fix bufferizesToMemoryRead for ml_program.global_store (PR #177387)

[Nathan Malimban]

================
@@ -50,3 +50,34 @@ func.func @raw_hazard() -> i64 {
   return %extracted : i64
 }
 
+// -----
+
+// CHECK-LABEL: memref.global "private" @state_tensor
+ml_program.global private mutable @"state_tensor"(dense<0.0> : tensor<4x75xf32>) : tensor<4x75xf32>
+
+// CHECK-LABEL: func.func @global_load_store_tensor
+func.func @global_load_store_tensor() -> tensor<4x75xf32> {
+  // CHECK-DAG:     %[[C0:.*]] = arith.constant 0
+  // CHECK-DAG:     %[[CST:.*]] = arith.constant 1.000000e+00
+  // CHECK-DAG:     %[[GLOB:.*]] = memref.get_global @state_tensor
+  // CHECK:         %[[VAL:.*]] = memref.load %[[GLOB]][%[[C0]], %[[C0]]]
+  // CHECK:         %[[ADD:.*]] = arith.addf %[[VAL]], %[[CST]]
+  // CHECK:         %[[ALLOC1:.*]] = memref.alloc() {alignment = 64 : i64}
+  // CHECK:         memref.copy %[[GLOB]], %[[ALLOC1]] 
+  // CHECK:         memref.store %[[ADD]], %[[ALLOC1]][%[[C0]], %[[C0]]] 
+  // CHECK:         %[[TENSOR:.*]] = bufferization.to_tensor %[[ALLOC1]] 
+  // CHECK:         %[[ALLOC2:.*]] = memref.alloc() {alignment = 64 : i64}
----------------
nmalimban wrote:

It seems that there are two `memref.alloc`s introduced: one to hold the result of the global load, another to hold the value to be stored back to global.

For the first, maybe it is because `GlobalLoadOpInterface` implements `isWritable` to return false, and that could cause `one-shot-bufferize` to introduce a copy to a new memref?

For the second, I see that `GlobalStoreOpInterface` implements `bufferize` by calling `createMemCpy`. So instead of copying the value directly from the first alloced memref to the global state, it allocs an intermediary memref and then copies from there? I'm not entirely sure how `one-shot-bufferize` goes about this.

https://github.com/llvm/llvm-project/pull/177387