Skip to content

Commit 156dfc4

Browse files
authored
Extend XeGPU fence op (#1070)
Adding fence op flush attribute and extending scope attribute along with SIMD lowering
1 parent 385cc49 commit 156dfc4

File tree

5 files changed

+128
-8
lines changed

5 files changed

+128
-8
lines changed
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
From cfb25d2f76986fb2566b3a26fd82c3f152d88a8a Mon Sep 17 00:00:00 2001
2+
From: Garra1980 <[email protected]>
3+
Date: Wed, 30 Apr 2025 11:24:51 -0500
4+
Subject: [PATCH] Extend xegpu op fence def
5+
6+
---
7+
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 24 +++++++++++++++++--
8+
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 5 ++--
9+
2 files changed, 25 insertions(+), 4 deletions(-)
10+
11+
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
12+
index f1bed70253ef..401043e590b1 100644
13+
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
14+
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
15+
@@ -138,14 +138,34 @@ def XeGPU_CacheHintAttr
16+
}
17+
18+
def XeGPU_FenceScopeWorkgroup: I32EnumAttrCase<"Workgroup", 0, "workgroup">;
19+
-def XeGPU_FenceScopeGPU: I32EnumAttrCase<"GPU", 1, "gpu">;
20+
+def XeGPU_FenceScopeLocal: I32EnumAttrCase<"Local", 1, "local">;
21+
+def XeGPU_FenceScopeTile: I32EnumAttrCase<"Tile", 2, "tile">;
22+
+def XeGPU_FenceScopeGPU: I32EnumAttrCase<"GPU", 3, "gpu">;
23+
+def XeGPU_FenceScopeSystem: I32EnumAttrCase<"System", 4, "system">;
24+
def XeGPU_FenceScope: I32EnumAttr<"FenceScope",
25+
"The enumeration for the scope of fence operation.",
26+
- [XeGPU_FenceScopeWorkgroup, XeGPU_FenceScopeGPU]> {
27+
+ [XeGPU_FenceScopeWorkgroup, XeGPU_FenceScopeLocal, XeGPU_FenceScopeTile, XeGPU_FenceScopeGPU,
28+
+ XeGPU_FenceScopeSystem]> {
29+
let genSpecializedAttr = 0;
30+
let cppNamespace = "::mlir::xegpu";
31+
}
32+
33+
+def XeGPU_FenceOpFlushNone: I32EnumAttrCase<"None", 0, "none">;
34+
+def XeGPU_FenceOpFlushEvict: I32EnumAttrCase<"Evict", 1, "evict">;
35+
+def XeGPU_FenceOpFlushInvalidate: I32EnumAttrCase<"Invalidate", 2, "invalidate">;
36+
+def XeGPU_FenceOpFlushDiscard: I32EnumAttrCase<"Discard", 3, "discard">;
37+
+def XeGPU_FenceOpFlush: I32EnumAttr<"FenceOpFlush",
38+
+ "The enumeration for the scope of fence cache flushing.",
39+
+ [XeGPU_FenceOpFlushNone, XeGPU_FenceOpFlushEvict, XeGPU_FenceOpFlushInvalidate, XeGPU_FenceOpFlushDiscard]> {
40+
+ let genSpecializedAttr = 0;
41+
+ let cppNamespace = "::mlir::xegpu";
42+
+}
43+
+def XeGPU_FenceOpFlushAttr:
44+
+ EnumAttr<XeGPU_Dialect, XeGPU_FenceOpFlush, "fence_cache_flushing"> {
45+
+ let summary = [{Describes the cache flushing for fence op.}];
46+
+ let assemblyFormat = "$value";
47+
+}
48+
+
49+
def XeGPU_FenceScopeAttr:
50+
EnumAttr<XeGPU_Dialect, XeGPU_FenceScope, "fence_scope"> {
51+
let summary = [{Describes the scope of fence.
52+
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
53+
index 9b001a78d6fe..01c5357e33ca 100644
54+
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
55+
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
56+
@@ -976,8 +976,9 @@ def XeGPU_FenceOp: XeGPU_Op<"fence", []> {
57+
within each workgroup. "GPU" means the scope would be across workgroups within the GPU.
58+
}];
59+
let arguments = (ins XeGPU_MemorySpaceAttr: $memory_kind,
60+
- XeGPU_FenceScopeAttr: $fence_scope);
61+
- let assemblyFormat = [{`memory_kind` `=` `` $memory_kind `,` `fence_scope` `=` `` $fence_scope attr-dict}];
62+
+ XeGPU_FenceScopeAttr: $fence_scope,
63+
+ OptionalAttr<XeGPU_FenceOpFlushAttr>: $fence_op_flush);
64+
+ let assemblyFormat = [{`memory_kind` `=` `` $memory_kind `,` `fence_scope` `=` `` $fence_scope (`,` `fence_op_flush` `=` `` $fence_op_flush^)? attr-dict}];
65+
let extraClassDeclaration = extraBaseClassDeclaration;
66+
}
67+
68+
--
69+
2.34.1

include/imex/Dialect/LLVMIR/XeVMOps.td

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -205,9 +205,10 @@ def XeVM_BlockStore2dOp : XeVM_Op<"blockstore2d">,
205205
def XeVM_MemoryScope : I32EnumAttr<"MemoryScope", "Memory scope for memory operations",
206206
[
207207
I32EnumAttrCase<"WORKGROUP", 0, "workgroup">,
208-
I32EnumAttrCase<"CLUSTER", 1, "cluster">,
209-
I32EnumAttrCase<"GPU", 2, "gpu">,
210-
I32EnumAttrCase<"SYSTEM", 3, "system">
208+
I32EnumAttrCase<"LOCAL", 1, "local">,
209+
I32EnumAttrCase<"TILE", 2, "tile">,
210+
I32EnumAttrCase<"GPU", 3, "gpu">,
211+
I32EnumAttrCase<"SYSTEM", 4, "system">
211212
]>{
212213
let cppNamespace = "imex::xevm";
213214
}

lib/Conversion/XeGPUToVC/LSCPatterns.cpp

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1317,16 +1317,17 @@ class FencePattern : public OpConversionPattern<FenceOp> {
13171317
ConversionPatternRewriter &rewriter) const override {
13181318
auto loc = op.getLoc();
13191319
auto pred = i1_val(1);
1320-
uint8_t fence_op, sfid, fence_scope;
1320+
uint8_t fence_op_flush, sfid, fence_scope;
13211321

1322-
enum lscFenceOp {
1322+
enum lscFenceOpFlush {
13231323
NONE = 0,
13241324
EVICT = 1,
13251325
INVALIDATE = 2,
13261326
DISCARD = 3,
13271327
CLEAN = 4,
13281328
FLUSHL3 = 5
13291329
};
1330+
// Fence scope values aligns with bspec but not all of them used
13301331
enum lscFenceScope {
13311332
GROUP = 0,
13321333
LOCAL = 1,
@@ -1338,8 +1339,7 @@ class FencePattern : public OpConversionPattern<FenceOp> {
13381339
};
13391340
enum lscSFID { UGM = 0, UGML = 1, TGM = 3, SLM = 4 };
13401341

1341-
// the design limits the fence_op to NONE
1342-
fence_op = lscFenceOp::NONE;
1342+
fence_op_flush = lscFenceOpFlush::NONE;
13431343
sfid = lscSFID::UGM;
13441344
fence_scope = lscFenceScope::GROUP;
13451345

@@ -1356,12 +1356,39 @@ class FencePattern : public OpConversionPattern<FenceOp> {
13561356
case xegpu::FenceScope::Workgroup:
13571357
fence_scope = lscFenceScope::GROUP;
13581358
break;
1359+
case xegpu::FenceScope::Local:
1360+
fence_scope = lscFenceScope::LOCAL;
1361+
break;
1362+
case xegpu::FenceScope::Tile:
1363+
fence_scope = lscFenceScope::TILE;
1364+
break;
13591365
case xegpu::FenceScope::GPU:
13601366
fence_scope = lscFenceScope::GPU;
13611367
break;
1368+
case xegpu::FenceScope::System:
1369+
fence_scope = lscFenceScope::SYSTEM;
1370+
break;
1371+
}
1372+
1373+
auto flushValue = op.getFenceOpFlush();
1374+
if (flushValue.has_value()) {
1375+
switch (flushValue.value()) {
1376+
case xegpu::FenceOpFlush::None:
1377+
fence_op_flush = lscFenceOpFlush::NONE;
1378+
break;
1379+
case xegpu::FenceOpFlush::Evict:
1380+
fence_op_flush = lscFenceOpFlush::EVICT;
1381+
break;
1382+
case xegpu::FenceOpFlush::Invalidate:
1383+
fence_op_flush = lscFenceOpFlush::INVALIDATE;
1384+
break;
1385+
case xegpu::FenceOpFlush::Discard:
1386+
fence_op_flush = lscFenceOpFlush::DISCARD;
1387+
break;
1388+
}
13621389
}
13631390

1364-
SmallVector<Value> args{pred, i8_val(sfid), i8_val(fence_op),
1391+
SmallVector<Value> args{pred, i8_val(sfid), i8_val(fence_op_flush),
13651392
i8_val(fence_scope)};
13661393

13671394
std::string funcName = "llvm.genx.lsc.fence.i1";

lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,9 +474,18 @@ class FenceToXeVMPattern : public OpConversionPattern<xegpu::FenceOp> {
474474
case xegpu::FenceScope::Workgroup:
475475
memScope = imex::xevm::MemoryScope::WORKGROUP;
476476
break;
477+
case xegpu::FenceScope::Local:
478+
memScope = imex::xevm::MemoryScope::LOCAL;
479+
break;
480+
case xegpu::FenceScope::Tile:
481+
memScope = imex::xevm::MemoryScope::TILE;
482+
break;
477483
case xegpu::FenceScope::GPU:
478484
memScope = imex::xevm::MemoryScope::GPU;
479485
break;
486+
case xegpu::FenceScope::System:
487+
memScope = imex::xevm::MemoryScope::SYSTEM;
488+
break;
480489
llvm_unreachable("Unknown XeGPU fence scope.");
481490
}
482491
imex::xevm::OclAddrSpace addrSpace{imex::xevm::OclAddrSpace::kGlobal};

test/Dialect/XeGPU/IR/barrier_ops.mlir

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,17 @@ func.func @fence() {
4343
xegpu.fence memory_kind = global, fence_scope = workgroup
4444
return
4545
}
46+
47+
// CHECK-LABEL: func @fence_1({{.*}}) {
48+
func.func @fence_1() {
49+
// CHECK: xegpu.fence memory_kind = global, fence_scope = tile
50+
xegpu.fence memory_kind = global, fence_scope = tile
51+
return
52+
}
53+
54+
// CHECK-LABEL: func @fence_2({{.*}}) {
55+
func.func @fence_2() {
56+
// CHECK: xegpu.fence memory_kind = global, fence_scope = tile, fence_op_flush = evict
57+
xegpu.fence memory_kind = global, fence_scope = tile, fence_op_flush = evict
58+
return
59+
}

0 commit comments

Comments
 (0)