Skip to content

Commit 4969ef9

Browse files
victor-edsvictor-eds
committed
[MLIR][GEN] Add GEN dialect
Add GEN dialect to represent operations on Intel GPUs. GEN will offer six initial operations: - `gen.local_id`: query a work-item's local id - `gen.work_group_id`: query the id of a work-item's work-group - `gen.work_group_size`: query the size of a work-item's work-group - `gen.num_work_groups`: query the number of work-groups - `gen.barrier`: work-group barrier - `gen.sub_group_shuffle`: sub-group shuffle Signed-off-by: Victor Perez <[email protected]>
1 parent 36230f9 commit 4969ef9

File tree

18 files changed

+421
-0
lines changed

18 files changed

+421
-0
lines changed

‎mlir/include/mlir/Dialect/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ add_subdirectory(ControlFlow)
1212
add_subdirectory(DLTI)
1313
add_subdirectory(EmitC)
1414
add_subdirectory(Func)
15+
add_subdirectory(GEN)
1516
add_subdirectory(GPU)
1617
add_subdirectory(Index)
1718
add_subdirectory(IRDL)

‎mlir/include/mlir/Dialect/GEN/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
add_subdirectory(IR)

‎mlir/include/mlir/Dialect/GEN/IR/CMakeLists.txt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
add_mlir_dialect(GENOps gen)
2+
add_mlir_doc(GENDialect GENDialect Dialects/ -gen-dialect-doc)
3+
add_mlir_doc(GENOps GENOps Dialects/ -gen-op-doc)
4+
5+
set(LLVM_TARGET_DEFINITIONS GENOps.td)
6+
mlir_tablegen(GENOpsEnums.h.inc -gen-enum-decls)
7+
mlir_tablegen(GENOpsEnums.cpp.inc -gen-enum-defs)
8+
add_public_tablegen_target(MLIRGENOpsEnumsIncGen)
9+
10+
set(LLVM_TARGET_DEFINITIONS GENAttrDefs.td)
11+
mlir_tablegen(GENOpsAttrDefs.h.inc -gen-attrdef-decls)
12+
mlir_tablegen(GENOpsAttrDefs.cpp.inc -gen-attrdef-defs)
13+
add_public_tablegen_target(MLIRGENOpsAttrDefsIncGen)

‎mlir/include/mlir/Dialect/GEN/IR/GENAttrDefs.td

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//===-- GENAttrDefs.td - GEN dialect attributes def. file --*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef GEN_ATTRDEFS
10+
#define GEN_ATTRDEFS
11+
12+
include "mlir/IR/EnumAttr.td"
13+
14+
/// Enum attribute of the different shuffle kinds.
15+
/// Based on https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_non_uniform_instructions
16+
def GEN_ShflKindAttr : I32EnumAttr<"ShflKind", "GEN shuffle kind",
17+
[
18+
I32EnumAttrCase<"XOR", 0, "xor">,
19+
I32EnumAttrCase<"UP", 1, "up">,
20+
I32EnumAttrCase<"DOWN", 2, "down">,
21+
I32EnumAttrCase<"IDX", 3, "idx">
22+
]> {
23+
let cppNamespace = "::mlir::GEN";
24+
}
25+
26+
#endif // GEN_ATTRDEFS

‎mlir/include/mlir/Dialect/GEN/IR/GENDialect.h

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
//===- GENDialect.h - MLIR GEN dialect --------------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file defines the GEN dialect in MLIR, containing Intel GEN operations.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef MLIR_DIALECT_GEN_IR_GENDIALECT_H
14+
#define MLIR_DIALECT_GEN_IR_GENDIALECT_H
15+
16+
#include "mlir/IR/Dialect.h"
17+
18+
#include "mlir/Dialect/GEN/IR/GENOpsDialect.h.inc"
19+
20+
#endif // MLIR_DIALECT_GEN_IR_GENDIALECT_H

‎mlir/include/mlir/Dialect/GEN/IR/GENDialect.td

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
//===-- GENDialect.td - GEN dialect op definition file -----*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef GEN_DIALECT
10+
#define GEN_DIALECT
11+
12+
include "mlir/IR/OpBase.td"
13+
14+
def GEN_Dialect : Dialect {
15+
let name = "gen";
16+
let cppNamespace = "::mlir::GEN";
17+
let summary = "The GEN dialect.";
18+
19+
let description = [{
20+
GEN is a dialect for representing operations on Intel GPUs.
21+
}];
22+
}
23+
24+
#endif // GEN_DIALECT

‎mlir/include/mlir/Dialect/GEN/IR/GENOps.h

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//===--- GENOps.h - GEN Dialect Operations ----------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef MLIR_DIALECT_GEN_IR_GENOPS_H
10+
#define MLIR_DIALECT_GEN_IR_GENOPS_H
11+
12+
#include "mlir/Bytecode/BytecodeOpInterface.h"
13+
#include "mlir/IR/OpDefinition.h"
14+
#include "mlir/Interfaces/SideEffectInterfaces.h"
15+
16+
#include "mlir/Dialect/GEN/IR/GENOpsEnums.h.inc"
17+
18+
#define GET_ATTRDEF_CLASSES
19+
#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.h.inc"
20+
21+
#include "mlir/Dialect/GEN/IR/GENTraits.h"
22+
23+
#define GET_OP_CLASSES
24+
#include "mlir/Dialect/GEN/IR/GENOps.h.inc"
25+
26+
#endif // MLIR_DIALECT_GEN_IR_GENOPS_H

‎mlir/include/mlir/Dialect/GEN/IR/GENOps.td

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
//===-- GENOps.td - GEN IR dialect op definition file ------*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This is the GEN IR operation definition file.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef GEN_OPS
14+
#define GEN_OPS
15+
16+
include "mlir/Dialect/GEN/IR/GENDialect.td"
17+
include "mlir/Dialect/GEN/IR/GENAttrDefs.td"
18+
include "mlir/IR/OpBase.td"
19+
include "mlir/IR/EnumAttr.td"
20+
include "mlir/Interfaces/SideEffectInterfaces.td"
21+
include "mlir/IR/OpAsmInterface.td"
22+
23+
//===----------------------------------------------------------------------===//
24+
// GEN op definitions
25+
//===----------------------------------------------------------------------===//
26+
27+
class GEN_Op<string mnemonic, list<Trait> traits = []> :
28+
Op<GEN_Dialect, mnemonic, traits>;
29+
30+
class GENOpTrait<string name, list<Trait> traits = [],
31+
code extraOpDeclaration = [{}],
32+
code extraOpDefinition = [{}]>
33+
: NativeOpTrait<name, traits, extraOpDeclaration, extraOpDefinition> {
34+
let cppNamespace = "::mlir::OpTrait::GEN";
35+
}
36+
37+
//===----------------------------------------------------------------------===//
38+
// ND-Range Operations
39+
//===----------------------------------------------------------------------===//
40+
41+
def GEN3DNDRange : GENOpTrait<"GEN3DNDRange">;
42+
43+
class GEN_3DNDRangeOp<string mnemonic, list<Trait> traits = []>
44+
: GEN_Op<mnemonic, [GEN3DNDRange, Pure] # traits>,
45+
Arguments<(ins I32:$dim)>,
46+
Results<(outs Index:$res)> {
47+
let assemblyFormat = "$dim attr-dict";
48+
}
49+
50+
def GEN_LocalIdOp : GEN_3DNDRangeOp<"local_id"> {
51+
let summary = "Query a work-item's local id.";
52+
let description = [{
53+
Query the work-item's position in its work-group, i.e., its local id, in a
54+
given dimension.
55+
```mlir
56+
%local_id = gen.local_id %dim
57+
```
58+
}];
59+
}
60+
61+
def GEN_WorkGroupIdOp : GEN_3DNDRangeOp<"work_group_id"> {
62+
let summary = "Query the id of a work-item's work-group.";
63+
let description = [{
64+
Query the id of a work-item's work-group in a given dimension.
65+
```mlir
66+
%work_group_id = gen.work_group_id %dim
67+
```
68+
}];
69+
}
70+
71+
def GEN_WorkGroupSizeOp : GEN_3DNDRangeOp<"work_group_size"> {
72+
let summary = "Query the work-group size.";
73+
let description = [{
74+
Query the work-item's work-group size in a given dimension.
75+
```mlir
76+
%work_group_size = gen.work_group_size %dim
77+
```
78+
}];
79+
}
80+
81+
def GEN_NumWorkGroupsOp : GEN_3DNDRangeOp<"num_work_groups"> {
82+
let summary = "Query the number of work-groups in the ND-range.";
83+
let description = [{
84+
Query the number of work-groups in the ND-range in a given dimension.
85+
```mlir
86+
%wg_number = gen.num_work_groups %dim
87+
```
88+
}];
89+
}
90+
91+
//===----------------------------------------------------------------------===//
92+
// Synchronization
93+
//===----------------------------------------------------------------------===//
94+
95+
def GEN_BarrierOp : GEN_Op<"barrier"> {
96+
let summary = "Work-group barrier";
97+
98+
string baseDescription = [{
99+
The `gen.barrier` operation performs a work-group barrier and ensures all
100+
outstanding memory transaction using local or global memory are complete.
101+
}];
102+
103+
let assemblyFormat = "attr-dict";
104+
}
105+
106+
def IntegerOrFloatType : AnyTypeOf<[AnySignlessInteger, AnyFloat]>;
107+
108+
def GEN_SubGroupShuffleOp : GEN_Op<"sub_group_shuffle", [
109+
TypesMatchWith<"result and value have the same type",
110+
"res", "value", "$_self">]>,
111+
Results<(outs IntegerOrFloatType:$res)>,
112+
Arguments<(ins IntegerOrFloatType:$value,
113+
I32:$mask,
114+
GEN_ShflKindAttr:$kind)> {
115+
let summary = "Sub-group shuffle";
116+
string baseDescription = [{
117+
The `gen.sub_group_shuffle` operation is invoked by different work items
118+
with different values, given by $value. Different work items have different
119+
sub-group local IDs. The shuffle kind, $kind, is given to determine how to
120+
calculate the associated sub-group local ID. It returns the associated
121+
$value for the work item with sub-group local ID equal to:
122+
- $kind == xor, the current invocation’s sub-group local ID xor’ed with $mask.
123+
- $kind == up, the current invocation’s sub-group local ID - $mask.
124+
- $kind == down, the current invocation’s sub-group local ID + $mask.
125+
- $kind == idx, the sub-group local ID $mask.
126+
}];
127+
128+
let assemblyFormat = [{
129+
$kind $value `,` $mask attr-dict `:` type($res)
130+
}];
131+
}
132+
133+
#endif // GEN_OPS

‎mlir/include/mlir/Dialect/GEN/IR/GENTraits.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
//===--- GENTraits.h - GEN Dialect Traits -----------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef MLIR_DIALECT_GEN_IR_GENTRAITS_H
10+
#define MLIR_DIALECT_GEN_IR_GENTRAITS_H
11+
12+
#include "mlir/IR/OpDefinition.h"
13+
14+
namespace mlir {
15+
namespace OpTrait {
16+
namespace GEN {
17+
namespace detail {
18+
LogicalResult verifyGEN3DNDRange(Operation *op);
19+
} // namespace detail
20+
21+
template <typename ConcreteType>
22+
class GEN3DNDRange : public TraitBase<ConcreteType, GEN3DNDRange> {
23+
public:
24+
static LogicalResult verifyTrait(Operation *op) {
25+
return detail::verifyGEN3DNDRange(op);
26+
}
27+
};
28+
} // namespace GEN
29+
} // namespace OpTrait
30+
} // namespace mlir
31+
32+
#endif // MLIR_DIALECT_GEN_IR_GENTRAITS_H

‎mlir/include/mlir/InitAllDialects.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
#include "mlir/Dialect/DLTI/DLTI.h"
3737
#include "mlir/Dialect/EmitC/IR/EmitC.h"
3838
#include "mlir/Dialect/Func/IR/FuncOps.h"
39+
#include "mlir/Dialect/GEN/IR/GENDialect.h"
3940
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
4041
#include "mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h"
4142
#include "mlir/Dialect/IRDL/IR/IRDL.h"
@@ -116,6 +117,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
116117
DLTIDialect,
117118
emitc::EmitCDialect,
118119
func::FuncDialect,
120+
GEN::GENDialect,
119121
gpu::GPUDialect,
120122
index::IndexDialect,
121123
irdl::IRDLDialect,

‎mlir/lib/Dialect/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ add_subdirectory(ControlFlow)
1212
add_subdirectory(DLTI)
1313
add_subdirectory(EmitC)
1414
add_subdirectory(Func)
15+
add_subdirectory(GEN)
1516
add_subdirectory(GPU)
1617
add_subdirectory(Index)
1718
add_subdirectory(IRDL)

‎mlir/lib/Dialect/GEN/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
add_subdirectory(IR)

‎mlir/lib/Dialect/GEN/IR/CMakeLists.txt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
add_mlir_dialect_library(MLIRGENDialect
2+
GENDialect.cpp
3+
GENOps.cpp
4+
GENTraits.cpp
5+
6+
ADDITIONAL_HEADER_DIRS
7+
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GEN
8+
9+
DEPENDS
10+
MLIRGENOpsIncGen
11+
MLIRGENOpsEnumsIncGen
12+
MLIRGENOpsAttrDefsIncGen
13+
14+
LINK_LIBS PUBLIC
15+
MLIRIR
16+
)

‎mlir/lib/Dialect/GEN/IR/GENDialect.cpp

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//===- GENDialect.cpp - MLIR GEN Dialect implementation -------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/Dialect/GEN/IR/GENDialect.h"
10+
11+
#include "mlir/Dialect/GEN/IR/GENOps.h"
12+
#include "mlir/IR/DialectImplementation.h"
13+
14+
using namespace mlir;
15+
using namespace mlir::GEN;
16+
17+
#include "mlir/Dialect/GEN/IR/GENOpsDialect.cpp.inc"
18+
#define GET_ATTRDEF_CLASSES
19+
#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.cpp.inc"
20+
21+
//===----------------------------------------------------------------------===//
22+
// GEN dialect.
23+
//===----------------------------------------------------------------------===//
24+
25+
void GENDialect::initialize() {
26+
addOperations<
27+
#define GET_OP_LIST
28+
#include "mlir/Dialect/GEN/IR/GENOps.cpp.inc"
29+
>();
30+
addAttributes<
31+
#define GET_ATTRDEF_LIST
32+
#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.cpp.inc"
33+
>();
34+
}

‎mlir/lib/Dialect/GEN/IR/GENOps.cpp

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
//===- GENOps.cpp - GEN dialect operations --------------------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/Dialect/GEN/IR/GENOps.h"
10+
11+
#include "mlir/IR/Builders.h"
12+
13+
//===----------------------------------------------------------------------===//
14+
// TableGen'd op method definitions
15+
//===----------------------------------------------------------------------===//
16+
17+
#define GET_OP_CLASSES
18+
#include "mlir/Dialect/GEN/IR/GENOps.cpp.inc"
19+
20+
//===----------------------------------------------------------------------===//
21+
// TableGen'd enum attribute definitions
22+
//===----------------------------------------------------------------------===//
23+
24+
#include "mlir/Dialect/GEN/IR/GENOpsEnums.cpp.inc"

0 commit comments

Comments
 (0)