Skip to content

Commit 2b4b4c0

Browse files
authored
[mlir][XeGPU] Add MemoryEffectsOpInterface for XeGPU memory related ops. (#125314)
`MemoryEffectsOpInterface` is needed for cleaning up the remaining (dead) operations after XeGPU distribution into SIMT.
1 parent 2f6e3df commit 2b4b4c0

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

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

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -276,7 +276,9 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", []> {
276276
}
277277

278278

279-
def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [AllElementTypesMatch<["value", "TensorDesc"]>]> {
279+
def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
280+
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemRead]>
281+
]> {
280282
let summary = "loads a n-D block from memory (represented by TensorDesc)"
281283
"to registers (represented by vector)";
282284
let description = [{
@@ -327,7 +329,9 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [AllElementTypesMatch<["value", "Tensor
327329
let hasVerifier = 1;
328330
}
329331

330-
def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [AllElementTypesMatch<["value", "TensorDesc"]>]> {
332+
def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
333+
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemWrite]>
334+
]> {
331335
let summary = "stores a n-D block register region back to memory, currently only supports 2D";
332336

333337
let description = [{
@@ -548,7 +552,9 @@ def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", []> {
548552
let hasVerifier = 1;
549553
}
550554

551-
def XeGPU_LoadGatherOp : XeGPU_Op<"load", [AllElementTypesMatch<["value", "TensorDesc"]>]> {
555+
def XeGPU_LoadGatherOp : XeGPU_Op<"load", [
556+
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemRead]>
557+
]> {
552558
let summary = "load a set of scattered data points from memory.";
553559

554560
let description = [{ It (aka. load) load data per each work-item. The output
@@ -618,7 +624,9 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [AllElementTypesMatch<["value", "Tenso
618624
let hasVerifier = 1;
619625
}
620626

621-
def XeGPU_StoreScatterOp : XeGPU_Op<"store", [AllElementTypesMatch<["value", "TensorDesc"]>]> {
627+
def XeGPU_StoreScatterOp : XeGPU_Op<"store", [
628+
AllElementTypesMatch<["value", "TensorDesc"]>, MemoryEffects<[MemWrite]>
629+
]> {
622630
let summary = "store data to scattered memory locations.";
623631
let description = [{ It (aka. store) stores data to scattered memory locations. The value is
624632
typically a 1D vector. But when the chunk size of the TensorDesc is larger than 1, it will be

0 commit comments

Comments
 (0)