[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.
+    // CHECK-LABEL: gpu.func @create_nd_tdesc_batch_subview(
+    // CHECK-SAME:  %[[SRC:.+]]: memref<4x64x64xf32>
+    gpu.func @create_nd_tdesc_batch_subview(%src: memref<4x64x64xf32>) -> vector<8xi32> {
+        // W = size[2] = 32, pitch = stride[1] = 64.
+        // CHECK: %[[W:.+]] = arith.trunci %{{.+}} : i64 to i32
+        // CHECK: %[[PITCH:.+]] = arith.trunci %{{.+}} : i64 to i32
+        // Batch row stride = stride[0] / pitch = 4096 / 64 = 64.
+        // CHECK: %[[ROWS0:.+]] = arith.constant 64 : i32
+        // base_height = size[1] + (size[0] - 1) * 64 = 32 + 3 * 64 = 224.
+        // CHECK: %[[H:.+]] = arith.trunci %{{.+}} : i64 to i32
+        // CHECK: %[[C1I32:.+]] = arith.constant 1 : i32
+        // CHECK: %[[BATCH:.+]] = arith.trunci %{{.+}} : i64 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: vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
+        %sub = memref.subview %src[0, 0, 0] [4, 32, 32] [1, 1, 1]
----------------
charithaintc wrote:

consider adding a dyamic shaped case for strides/memref sizes. 

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