[Mlir-commits] [mlir] [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (PR #215711)

[Jianhui Li]

================
@@ -232,28 +235,52 @@ class CreateNdDescToXeVMPattern
     Value baseShapeH;
 
     // Source can be a memref or a pointer (ui64, ui32, i64 or i32).
-    SmallVector<OpFoldResult> mixedSizes = op.getMixedSizes();
-    SmallVector<OpFoldResult> mixedStrides = op.getMixedStrides();
-    // Descriptor shape is expected to be 2D.
-    int64_t rank = mixedSizes.size();
     auto sourceTy = source.getType();
     auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
+
+    // For a memref source, shape/strides come from the memref (dynamic dims are
+    // recovered from runtime metadata); an integer source carries them as
+    // explicit op operands.
+    SmallVector<OpFoldResult> mixedSizes;
+    SmallVector<OpFoldResult> mixedStrides;
     // If source is a memref, we need to extract the aligned pointer as index.
     // Pointer type is passed as i32 or i64 by type converter.
     if (sourceMemrefTy) {
       if (!sourceMemrefTy.hasRank()) {
         return rewriter.notifyMatchFailure(op, "Expected ranked Memref.");
       }
+      SmallVector<int64_t> staticStrides;
+      int64_t staticOffset;
+      if (failed(
+              sourceMemrefTy.getStridesAndOffset(staticStrides, staticOffset)))
+        return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
+      // A fully static memref yields constants directly; a dynamic one recovers
+      // its dynamic dims from runtime metadata.
+      bool allStatic = sourceMemrefTy.hasStaticShape() &&
+                       llvm::none_of(staticStrides, ShapedType::isDynamic);
+      if (allStatic) {
+        mixedSizes = op.getMixedSizes();
+        mixedStrides = op.getMixedStrides();
+      } else {
+        auto srcMeta =
+            memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
+        mixedSizes = srcMeta.getConstifiedMixedSizes();
+        mixedStrides = srcMeta.getConstifiedMixedStrides();
+      }
       // Access adaptor after failure check to avoid rolling back generated code
       // for materialization cast.
       baseAddr = adaptor.getSource();
     } else {
+      mixedSizes = op.getMixedSizes();
+      mixedStrides = op.getMixedStrides();
       baseAddr = adaptor.getSource();
       if (baseAddr.getType() != i64Ty) {
         // Pointer type may be i32. Cast to i64 if needed.
         baseAddr = arith::ExtUIOp::create(rewriter, loc, i64Ty, baseAddr);
       }
     }
+    // Descriptor shape rank.
+    int64_t rank = mixedSizes.size();
----------------
Jianhui-Li wrote:

added a check upfront to verify the assumption tensor_desc rank is equal to memref rank. 

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