[Mlir-commits] [mlir] [mlir][XeGPU][XeGPUUnroll] Support new syntax with offsets moved to load_nd/store_nd/prefetch_nd (PR #160323)

[Dmitry Chigarev]

================
@@ -121,53 +121,79 @@ struct UnrollPattern : public OpRewritePattern<SourceOp> {
   xegpu::UnrollOptions options;
 };
 
+// Generic helper function for unrolling operations with offsets.
+//
+// Iterates over tile offsets within the tensor descriptor shape and calls
+// the provided createOp function for each computed offset. This is used by
+// operations like LoadNd, StoreNd, CreateNdDesc, and PrefetchNd when they
+// have explicit offsets that need to be adjusted for each unrolled tile.
+SmallVector<Value> computeUnrolledOffsets(
+    SmallVector<OpFoldResult> mixedOffsets, xegpu::TensorDescType tdescTy,
+    ArrayRef<int64_t> targetShape,
+    const std::function<Value(SmallVector<OpFoldResult>)> &createOp,
+    Location loc, PatternRewriter &rewriter) {
+  int64_t rank = tdescTy.getRank();
+  ArrayRef<int64_t> shape = tdescTy.getShape();
----------------
dchigarev wrote:

the logic of this function was copied from [`UnrollCreateNdOp`](https://github.com/llvm/llvm-project/blob/0ebcc9d0bcd53d8f02f358e52d40a7ab25cdaa4e/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp#L139-L171), it is now used in UnrollCreateNdOp/LoadNd/StoreNd/PrefetchNd.

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