[MLIR][XeGPU] Allow load/store/prefetch uses [memref+offset] instead of tdesc #150576

Jianhui-Li · 2025-07-25T05:32:55Z

Add variant of load/store/prefetch to allow offset. The new xegpu.load variant accepts memref+offset, and the existing tdesc operand will be removed in the future PR.

The semantics are combination of "creating scattered_tdesc + xegpu.load with scattered_tdesc". The current xegpu.load accepts tdesc operand, which encapsulates "memref+offset". This PR "fold" "memref+offset" directly to xegpu.load replacing "tdesc". Create_tdesc will be removed as scatter_tdesc only contains base address after offsets being taken away, so there is no point to keep it.

    // wi level code example
    %2 = xegpu.load %src[%offsets], %mask <{chunk_size = 2}> : ui64,  vector<1xindex>, vector<1xi1> -> vector<2xf32>
    xegpu.store %val, %src[%offsets], %mask: vector<1xf16>, memref<?xf16>, vector<1xindex>, vector<1xi1>
    xegpu.prefetch %src[%0] : ui64, vector<1xindex>

chencha3

LGTM.

chencha3 · 2025-07-25T20:14:07Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

  let genVerifyDecl = 1;
 }

+def XeGPU_TensorDesc_or_MemRef : AnyTypeOf<[XeGPU_TensorDesc,Non0RankedMemRefOf<[XeGPU_ScalarType]>, UI64]>;


nit: keep a consistent naming convention: XeGPU_TensorDescOrMemRef

chencha3 · 2025-07-25T20:17:39Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

 include "mlir/Interfaces/ShapedOpInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
+include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"


I am not sure whether this header include is necessary, So far, I didn't see the changes requiring this. Maybe It can be removed.

chencha3 · 2025-07-25T20:23:06Z

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

+
+  auto maskShape = getShapeOf(maskTy);
+  auto valueShape = getShapeOf(valueTy);
+  auto memShape = getShapeOf(memTy);


It seems memShape is not used. Maybe it can be removed.

adam-smnk

There's a bit too much happening here.
Could this PR be split?

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

charithaintc

overall LGTM. please address the comments.

charithaintc · 2025-07-28T18:16:16Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

+      return getSource().getType();
+    }
+
+    Value getTensorDesc() {


why this method is needed?

to minimize the change during the transition.

nit: this is bit risky. because the caller might expect a tensor_desc but this may return a memref or I64. I would add a TODO note.

Since it returns one specific variant of the source, I'd also suggest it return a TypedValue<xegpu::TensorDescType>

charithaintc · 2025-07-28T18:17:23Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

+       return getSource();
+    }
+
    xegpu::TensorDescType getTensorDescType() {


here dyn_cast is involved. Maybe better to rename the function to TryGetTensorDesc to avoid confusion?

The function name needs to keep the same to minimize the change.

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

charithaintc · 2025-07-28T18:58:50Z

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

+}
+
+static LogicalResult
+isValidGatherScatterRawptrParams(Type maskTy, VectorType valueTy,


these two functions are very similar. I think we can reuse/refactor isValidGatherScatterParams to achive this. I don't see a need to define 2 new functions.

If that is hard to do, at least consider moving common logic to a helper and reuse the helper.

refactored to one function.

charithaintc · 2025-07-28T20:26:05Z

mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp

    xegpu::TensorDescType tdescTy = op.getTensorDescType();

-    if (!tdescTy.isScattered())
+    if (!tdescTy || !tdescTy.isScattered())


nit: Looks like the new version is supported by this pattern? Maybe add a TODO note will help here.

charithaintc · 2025-07-28T20:34:25Z

mlir/test/Dialect/XeGPU/invalid.mlir

+}
+
+// -----
+func.func @store_scatter_offset_sg(%src: memref<?xf16>) {


is this a wi test case like the one above? mask and offsets are just 1 x type

mlir/test/Dialect/XeGPU/ops.mlir

Jianhui-Li · 2025-07-29T17:42:17Z

@adam-smnk @charithaintc Thanks for the detailed review! I have added documentation and updated the PR description. Also changed the type name and refactor the verifier.

charithaintc

LGTM. please address the minor comments if possible.

charithaintc · 2025-07-29T17:48:00Z

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

+
+  if (tdescTy) {
+    if (!tdescTy.isScattered())
+      return emitOpError("Expects a scattered TensorDesc.\n");
+  } else {
+    if (getRankOf(getSource()) > 1)
+      return emitOpError(
+          "Expecting the source is a 1D memref or pointer (uint64_t).");
+  }


Suggested change

if (tdescTy) {

if (!tdescTy.isScattered())

return emitOpError("Expects a scattered TensorDesc.\n");

} else {

if (getRankOf(getSource()) > 1)

return emitOpError(

"Expecting the source is a 1D memref or pointer (uint64_t).");

}

if (tdescTy && !tdescTy.isScattered())

return emitOpError("Expects a scattered TensorDesc.\n");

if (getRankOf(getSource()) > 1)

return emitOpError(

"Expecting the source is a 1D memref or pointer (uint64_t).");

The suggested code doesn't work since the Tdesc can be 2d.

Suggested change

if (tdescTy) {

if (!tdescTy.isScattered())

return emitOpError("Expects a scattered TensorDesc.\n");

} else {

if (getRankOf(getSource()) > 1)

return emitOpError(

"Expecting the source is a 1D memref or pointer (uint64_t).");

}

if (tdescTy && !tdescTy.isScattered())

return emitOpError("Expects a scattered TensorDesc.\n");

if (!tdescTy && getRankOf(getSource()) > 1)

return emitOpError(

"Expecting the source is a 1D memref or pointer (uint64_t).");

charithaintc · 2025-07-29T17:49:17Z

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

+  if (tdescTy) {
+    if (!tdescTy.isScattered())
+      return emitOpError("Expects a scattered TensorDesc.\n");
+  } else {
+    if (getRankOf(getSource()) > 1)
+      return emitOpError(
+          "Expecting the source is a 1D memref or pointer (uint64_t).");
+  }


check the suggested change above. in summary, we should try to combine conditions and return early rather than nesting if conditions.

charithaintc · 2025-07-29T17:49:48Z

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

+  if (tdescTy) {
+    if (!tdescTy.isScattered())
+      return emitOpError("Expects a scattered TensorDesc.\n");
+  } else {
+    if (getRankOf(getDest()) > 1)
+      return emitOpError(
+          "Expecting the dest is a 1D memref or pointer (uint64_t).");
+  }


check suggested code change.

charithaintc · 2025-07-29T17:51:35Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

+      return getSource().getType();
+    }
+
+    Value getTensorDesc() {


nit: this is bit risky. because the caller might expect a tensor_desc but this may return a memref or I64. I would add a TODO note.

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

adam-smnk · 2025-07-30T12:21:58Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

+      return getSource().getType();
+    }
+
+    Value getTensorDesc() {


Since it returns one specific variant of the source, I'd also suggest it return a TypedValue<xegpu::TensorDescType>

mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

adam-smnk

LGTM % previous nits

Jianhui-Li added 13 commits July 17, 2025 23:21

add optional offsets to nd load/store/prefetch

1373ffa

git-clang-format

30ff640

Merge branch 'main' into load-nd-with-offsets-2

037f2a2

add optional offsets to load_gather

efd1661

add optional offsets to nd load/store/prefetch

3578c1b

add optional offsets to load_gather

59f7ea9

add offsets to load

abc84c7

add chunk_size and use XeGPU_offsetType

80b4462

fix conflict

0b91942

add tests

769bf19

add invalid tests

4553746

Merge branch 'main' into load-with-offsets

e6c4db3

small fixes

1249794

adam-smnk requested review from adam-smnk, charithaintc and chencha3 July 25, 2025 15:57

chencha3 approved these changes Jul 25, 2025

View reviewed changes

address comments

5cfb24b

adam-smnk reviewed Jul 28, 2025

View reviewed changes

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td Outdated Show resolved Hide resolved

adam-smnk reviewed Jul 28, 2025

View reviewed changes

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp Outdated Show resolved Hide resolved

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td Outdated Show resolved Hide resolved

charithaintc reviewed Jul 28, 2025

View reviewed changes

Jianhui-Li changed the title ~~[MLIR][XeGPU] Add offsets to load/store/prefetch~~ [MLIR][XeGPU] Allow load/store/prefetch uses [memref+offset] instead of tdesc Jul 29, 2025

Jianhui-Li added 2 commits July 29, 2025 17:33

address feedback

5940d19

git-clang-format

da7142a

charithaintc self-requested a review July 29, 2025 17:45

charithaintc approved these changes Jul 29, 2025

View reviewed changes

minor polish

8b99ecc

adam-smnk reviewed Jul 30, 2025

View reviewed changes

address comments

04306ca

Merge branch 'main' into load-with-offsets

871e4d4

adam-smnk approved these changes Jul 30, 2025

View reviewed changes

add more invalid tests

bbd6530

Jianhui-Li merged commit e6f360b into llvm:main Jul 30, 2025
9 checks passed

[MLIR][XeGPU] Allow load/store/prefetch uses [memref+offset] instead of tdesc #150576

[MLIR][XeGPU] Allow load/store/prefetch uses [memref+offset] instead of tdesc #150576

Uh oh!

Conversation

Jianhui-Li commented Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chencha3 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adam-smnk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Jianhui-Li commented Jul 29, 2025

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Jianhui-Li commented Jul 25, 2025 •

edited

Loading