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

[Tuomas Kärnä]

================
@@ -433,6 +433,62 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: @insert_prefetch_dynamic
+func.func @insert_prefetch_dynamic(%arg0: memref<128x?xf16>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  // CHECK:    %[[K:.*]] = memref.dim %arg0
+  %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+  // CHECK:    %[[OUTER:.*]] = memref.subview %arg0[0, 0]
+  %outer = memref.subview %arg0[0, 0] [128, %k] [1, 1]
+      : memref<128x?xf16> to memref<128x?xf16, strided<[?, 1]>>
+  // CHECK:    %[[INIT_END:.*]] = arith.minsi {{.*}}, %[[K]]
+  // CHECK:    scf.for %[[INIT_IV:.*]] = %{{.*}} to %[[INIT_END]] step %{{.*}} {
+  // CHECK:      %[[INIT_SUB:.*]] = memref.subview %[[OUTER]][0, %[[INIT_IV]]]
+  // CHECK:      %[[INIT_DESC:.*]] = xegpu.create_nd_tdesc %[[INIT_SUB]]
+  // CHECK:      xegpu.prefetch_nd %[[INIT_DESC]][0, 0]
+  // CHECK:    }
+  // CHECK:    scf.for %[[IV:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
+  // CHECK:      %[[FUTURE:.*]] = arith.addi %[[IV]]
+  // CHECK:      %[[IN_BOUNDS:.*]] = arith.cmpi slt, %[[FUTURE]], %[[K]]
+  // CHECK:      scf.if %[[IN_BOUNDS]] {
+  // CHECK:        %[[MAIN_SUB:.*]] = memref.subview %[[OUTER]][0, %[[FUTURE]]]
+  // CHECK:        %[[MAIN_DESC:.*]] = xegpu.create_nd_tdesc %[[MAIN_SUB]]
+  // CHECK:        xegpu.prefetch_nd %[[MAIN_DESC]][0, 0]
+  // CHECK:      }
+  // CHECK:      xegpu.load_nd
+  // CHECK:    }
+  scf.for %iv = %c0 to %k step %c16 {
+    %remaining = arith.subi %k, %iv : index
+    %size = arith.minsi %remaining, %c16 : index
+    %tile = memref.subview %outer[0, %iv] [128, %size] [1, 1]
+        : memref<128x?xf16, strided<[?, 1]>>
+          to memref<128x?xf16, strided<[?, 1], offset: ?>>
+    %desc = xegpu.create_nd_tdesc %tile
+        : memref<128x?xf16, strided<[?, 1], offset: ?>>
+          -> !xegpu.tensor_desc<128x16xf16>
----------------
tkarna wrote:

On a high level, it's unclear to me why in the case of dynamic shapes we end up with this IR form. Why can't we just create the descriptor outside the loop with 0 offset and then compute the right offset in the loop with `%iv` and use that in the load op?

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