Skip to content

Commit daebe5c

Browse files
chencha3joker-eph
andauthored
[MLIR][XeGPU] Adding XeGPU 2d block operators (llvm#84692)
Hi @joker-eph, This PR adds XeGPU 2D block operators. It contains: 1. `TensorDescType` and `TensorDescAttr` definitions 2. `MemoryScopeAttr` and `CacheHintAttr` definitions which are used by `TensorDescAttr`. 3. `CreateNdDescOp`, `PrefetchNdOp`, `LoadNdOp`, and `StoreNdOp` definitions, and their corresponding testcases for illustration. --------- Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
1 parent a9f78a3 commit daebe5c

File tree

8 files changed

+767
-12
lines changed

8 files changed

+767
-12
lines changed

mlir/include/mlir/Dialect/XeGPU/IR/XeGPU.h

+6-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,12 @@
99
#ifndef MLIR_DIALECT_XEGPU_IR_XEGPU_H
1010
#define MLIR_DIALECT_XEGPU_IR_XEGPU_H
1111

12-
#include <mlir/IR/Dialect.h>
12+
#include "mlir/Bytecode/BytecodeOpInterface.h"
13+
#include "mlir/IR/BuiltinTypes.h"
14+
#include "mlir/IR/Dialect.h"
15+
#include "mlir/Interfaces/ShapedOpInterfaces.h"
16+
#include "mlir/Interfaces/SideEffectInterfaces.h"
17+
#include "mlir/Interfaces/ViewLikeInterface.h"
1318

1419
namespace mlir {
1520
namespace xegpu {

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

+61
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,72 @@
1010
#define MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD
1111

1212
include "mlir/Dialect/XeGPU/IR/XeGPUDialect.td"
13+
include "mlir/IR/EnumAttr.td"
1314

1415
class XeGPUAttr<string name, string attrMnemonic, list<Trait> traits = [],
1516
string baseCppClass = "::mlir::Attribute">
1617
: AttrDef<XeGPU_Dialect, name, traits, baseCppClass> {
1718
let mnemonic = attrMnemonic;
1819
}
1920

21+
def XeGPU_TensorDescAttr: XeGPUAttr<"TensorDesc", "tdesc_attr"> {
22+
let parameters = (ins
23+
OptionalParameter<"MemoryScopeAttr">: $memory_scope,
24+
OptionalParameter<"IntegerAttr", "1">: $array_length,
25+
OptionalParameter<"BoolAttr", "true">: $boundary_check
26+
);
27+
28+
let builders = [
29+
AttrBuilder<(ins
30+
CArg<"xegpu::MemoryScope", "xegpu::MemoryScope::Global">:$memory_scope,
31+
CArg<"int", "1">:$array_length,
32+
CArg<"bool", "true">: $boundary_check
33+
)>
34+
];
35+
36+
let assemblyFormat = "`<` struct(params) `>`";
37+
}
38+
39+
//===----------------------------------------------------------------------===//
40+
// XeGPU Memory Scope Enums.
41+
//===----------------------------------------------------------------------===//
42+
def XeGPU_MemoryScopeGlobal: I32EnumAttrCase<"Global", 0, "global">;
43+
def XeGPU_MemoryScopeShared: I32EnumAttrCase<"SLM", 1, "slm">;
44+
def XeGPU_MemoryScope: I32EnumAttr<"MemoryScope",
45+
"The address space of the memory the tensor descritor is created for",
46+
[XeGPU_MemoryScopeGlobal, XeGPU_MemoryScopeShared]> {
47+
let genSpecializedAttr = 0;
48+
let cppNamespace = "::mlir::xegpu";
49+
}
50+
51+
def XeGPU_MemoryScopeAttr:
52+
EnumAttr<XeGPU_Dialect, XeGPU_MemoryScope, "memory_scope"> {
53+
let assemblyFormat = "$value";
54+
}
55+
56+
//===----------------------------------------------------------------------===//
57+
// XeGPU Cache Enums.
58+
//===----------------------------------------------------------------------===//
59+
def XeGPU_CachePolicyCached: I32EnumAttrCase<"CACHED", 0, "cached">; // valid for read and write
60+
def XeGPU_CachePolicyUncached: I32EnumAttrCase<"UNCACHED", 1, "uncached">; // valid for read and write
61+
def XeGPU_CachePolicyStreaming: I32EnumAttrCase<"STREAMING", 2, "streaming">; // valid for read only
62+
def XeGPU_CachePolicyInvalid: I32EnumAttrCase<"READ_INVALIDATE", 3, "read_invalidate">; // valid for read only
63+
def XeGPU_CachePolicyWriteBack: I32EnumAttrCase<"WRITE_BACK", 4, "write_back">; // valid for write only
64+
def XeGPU_CachePolicyWriteThrough: I32EnumAttrCase<"WRITE_THROUGH", 5, "write_through">; // valid for write only
65+
66+
def XeGPU_CachePolicyEnums : I32EnumAttr<"CachePolicy", "Cache policy",
67+
[XeGPU_CachePolicyCached, XeGPU_CachePolicyUncached,
68+
XeGPU_CachePolicyStreaming, XeGPU_CachePolicyInvalid,
69+
XeGPU_CachePolicyWriteBack, XeGPU_CachePolicyWriteThrough]> {
70+
let genSpecializedAttr = 0;
71+
let cppNamespace = "::mlir::xegpu";
72+
}
73+
74+
def XeGPU_CacheHintAttr
75+
: EnumAttr<XeGPU_Dialect, XeGPU_CachePolicyEnums, "cache_hint"> {
76+
let assemblyFormat = "`<` $value `>`";
77+
}
78+
79+
80+
2081
#endif // MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@ def XeGPU_Dialect : Dialect {
2323
the lower-level GPU compiler.
2424
}];
2525

26-
// let useDefaultTypePrinterParser = true;
27-
// let useDefaultAttributePrinterParser = true;
26+
let useDefaultTypePrinterParser = true;
27+
let useDefaultAttributePrinterParser = true;
2828
}
2929

3030
#endif // MLIR_DIALECT_XEGPU_IR_XEGPUDIALECT_TD

0 commit comments

Comments
 (0)