[Mlir-commits] [mlir] [mlir][xegpu] Fix high-D nd base_height for non-packed plane strides (PR #225556)

[Charitha Saumya]

================
@@ -60,24 +60,58 @@ gpu.module @create_nd_tdesc {
         gpu.return
     }
 
-    // Batched (>2D): base_height spans all planes; slot 5 = batch row stride.
+    // Batched (>2D): base_height is the row extent the batch dim reaches,
+    // size[1] + (size[0] - 1) * batch_row_stride; slot 5 = batch row stride.
     // CHECK-LABEL: gpu.func @create_nd_tdesc_batch_dyn(
     // CHECK-SAME:  %[[SRC:.+]]: memref<?x?x?xf16>
     gpu.func @create_nd_tdesc_batch_dyn(%src: memref<?x?x?xf16>) -> vector<8xi32> {
         // CHECK: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
         // CHECK: %[[W:.+]] = arith.index_cast %[[SIZES]]#2 : index to i32
+        // CHECK: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+        // CHECK: %[[LS0:.+]] = arith.index_cast %[[STRIDES]]#0 : index to i32
+        // CHECK: %[[ROWS0:.+]] = arith.divui %[[LS0]], %[[PITCH]] : i32
         // CHECK: %[[H:.+]] = arith.index_cast %[[SIZES]]#1 : index to i32
+        // CHECK: %[[C1I32:.+]] = arith.constant 1 : i32
         // CHECK: %[[BATCH:.+]] = arith.index_cast %[[SIZES]]#0 : index to i32
-        // CHECK: %[[FLAT_H:.+]] = arith.muli %[[H]], %[[BATCH]] : i32
-        // CHECK: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+        // CHECK: %[[BATCHM1:.+]] = arith.subi %[[BATCH]], %[[C1I32]] : i32
+        // CHECK: %[[BATCH_ROWS:.+]] = arith.muli %[[BATCHM1]], %[[ROWS0]] : i32
+        // CHECK: %[[FLAT_H:.+]] = arith.addi %[[H]], %[[BATCH_ROWS]] : i32
         // CHECK: %[[P2:.+]] = vector.insert %[[W]], %{{.+}} [2] : i32 into vector<8xi32>
         // CHECK: %[[P3:.+]] = vector.insert %[[FLAT_H]], %[[P2]] [3] : i32 into vector<8xi32>
         // CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
-        // CHECK: %[[LS0:.+]] = arith.index_cast %[[STRIDES]]#0 : index to i32
-        // CHECK: %[[ROWS0:.+]] = arith.divui %[[LS0]], %[[PITCH]] : i32
         // CHECK: vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
         %t = xegpu.create_nd_tdesc %src : memref<?x?x?xf16> -> !xegpu.tensor_desc<1x8x16xf16>
         %c = builtin.unrealized_conversion_cast %t : !xegpu.tensor_desc<1x8x16xf16> to vector<8xi32>
         gpu.return %c : vector<8xi32>
     }
+
+    // Batched (>2D) over a subview of a wider source: the planes are 64 rows
+    // apart but only 32 rows tall, so base_height is the row extent
+    // 32 + 3 * 64 = 224.
----------------
charithaintc wrote:

why not use 32 as the height. 224 value does not have a real meaning here, we still have to jump over 32 rows to reach the second plane if we have non zero offsets. 

Consider adding a case with nonzero offsets and check if the formula still works for that case. 

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