[Mlir-commits] [mlir] [xegpu][transform] Improve `insert_prefetch` to support loads from dynamic shaped memrefs. (PR #228258)

[Charitha Saumya]

================
@@ -101,3 +101,142 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+func.func @insert_prefetch_loop_local_offset(%arg0: memref<128x?xf16>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+  scf.for %iv = %c0 to %k step %c16 {
+    %remaining = arith.subi %k, %iv : index
+    %size = arith.minsi %remaining, %c16 : index
+    %tile = memref.subview %arg0[0, %iv] [128, %size] [1, 1]
+        : memref<128x?xf16> to memref<128x?xf16, strided<[?, 1], offset: ?>>
+    %desc = xegpu.create_nd_tdesc %tile
+        : memref<128x?xf16, strided<[?, 1], offset: ?>>
+          -> !xegpu.tensor_desc<128x16xf16>
+    %offset = arith.addi %iv, %c0 : index
+    %load = xegpu.load_nd %desc[0, %offset]
+        : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>
+  }
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    // expected-error at below {{loop-local prefetch requires zero load offsets}}
+    %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+func.func @insert_prefetch_no_iv(%arg0: memref<128x128xf16>) {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %c128 = arith.constant 128 : index
+  %desc = xegpu.create_nd_tdesc %arg0
+      : memref<128x128xf16> -> !xegpu.tensor_desc<16x16xf16>
+  scf.for %iv = %c0 to %c128 step %c16 {
+    %load = xegpu.load_nd %desc[0, 0]
+        : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+  }
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    // expected-error at below {{prefetch load offsets must directly use the loop induction variable}}
+    %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+func.func @insert_prefetch_derived_load_offset(%arg0: memref<128x128xf16>) {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %c128 = arith.constant 128 : index
+  %desc = xegpu.create_nd_tdesc %arg0
+      : memref<128x128xf16> -> !xegpu.tensor_desc<16x16xf16>
+  scf.for %iv = %c0 to %c128 step %c16 {
+    %offset = arith.addi %iv, %c0 : index
+    %load = xegpu.load_nd %desc[%offset, %iv]
----------------
charithaintc wrote:

hmm, I agree this needs to be supported.

but it seems this fails currently on main. I am thinking about adding general support for affine offsets.

```
./debug_dir/prefetch_offset_test.mlir:5:11: error: operand #1 does not dominate this use
  %desc = xegpu.create_nd_tdesc %arg0
          ^
./debug_dir/prefetch_offset_test.mlir:5:11: note: see current operation: "xegpu.prefetch_nd"(%4, %9, %0) <{const_offsets = array<i64: -9223372036854775808, -9223372036854775808>, l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>}> : (!xegpu.tensor_desc<16x16xf16>, index, index) -> ()
./debug_dir/prefetch_offset_test.mlir:8:15: note: operand defined here (op in a child region)
    %offset = arith.addi %iv, %c0 : index
```


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