[MLIR][GEN] Add GEN dialect

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -12,6 +12,7 @@ add_subdirectory(ControlFlow)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(Func)
+add_subdirectory(GEN)
 add_subdirectory(GPU)
 add_subdirectory(Index)
 add_subdirectory(IRDL)

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

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,13 @@
+add_mlir_dialect(GENOps gen)
+add_mlir_doc(GENDialect GENDialect Dialects/ -gen-dialect-doc)
+add_mlir_doc(GENOps GENOps Dialects/ -gen-op-doc)
+
+set(LLVM_TARGET_DEFINITIONS GENOps.td)
+mlir_tablegen(GENOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(GENOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRGENOpsEnumsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS GENAttrDefs.td)
+mlir_tablegen(GENOpsAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(GENOpsAttrDefs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRGENOpsAttrDefsIncGen)

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

@@ -0,0 +1,26 @@
+//===-- GENAttrDefs.td - GEN dialect attributes def. file --*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GEN_ATTRDEFS
+#define GEN_ATTRDEFS
+
+include "mlir/IR/EnumAttr.td"
+
+/// Enum attribute of the different shuffle kinds.
+/// Based on https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_non_uniform_instructions
+def GEN_ShflKindAttr : I32EnumAttr<"ShflKind", "GEN shuffle kind",
+  [
+    I32EnumAttrCase<"XOR",  0, "xor">,
+    I32EnumAttrCase<"UP",   1, "up">,
+    I32EnumAttrCase<"DOWN", 2, "down">,
+    I32EnumAttrCase<"IDX",  3, "idx">
+  ]> {
+  let cppNamespace = "::mlir::GEN";
+}
+
+#endif // GEN_ATTRDEFS

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

@@ -0,0 +1,20 @@
+//===- GENDialect.h - MLIR GEN dialect --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the GEN dialect in MLIR, containing Intel GEN operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_GEN_IR_GENDIALECT_H
+#define MLIR_DIALECT_GEN_IR_GENDIALECT_H
+
+#include "mlir/IR/Dialect.h"
+
+#include "mlir/Dialect/GEN/IR/GENOpsDialect.h.inc"
+
+#endif // MLIR_DIALECT_GEN_IR_GENDIALECT_H

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

@@ -0,0 +1,24 @@
+//===-- GENDialect.td - GEN dialect op definition file -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GEN_DIALECT
+#define GEN_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def GEN_Dialect : Dialect {
+  let name = "gen";
+  let cppNamespace = "::mlir::GEN";
+  let summary = "The GEN dialect.";
+
+  let description = [{
+    GEN is a dialect for representing operations on Intel GPUs.
+  }];
+}
+
+#endif // GEN_DIALECT

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

@@ -0,0 +1,24 @@
+//===--- GENOps.h - GEN Dialect Operations ----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_GEN_IR_GENOPS_H
+#define MLIR_DIALECT_GEN_IR_GENOPS_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "mlir/Dialect/GEN/IR/GENOpsEnums.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.h.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/GEN/IR/GENOps.h.inc"
+
+#endif // MLIR_DIALECT_GEN_IR_GENOPS_H

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

@@ -0,0 +1,166 @@
+//===-- GENOps.td - GEN IR dialect op definition file ------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the GEN IR operation definition file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GEN_OPS
+#define GEN_OPS
+
+include "mlir/Dialect/GEN/IR/GENDialect.td"
+include "mlir/Dialect/GEN/IR/GENAttrDefs.td"
+include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpAsmInterface.td"
+
+//===----------------------------------------------------------------------===//
+// GEN op definitions
+//===----------------------------------------------------------------------===//
+
+class GEN_Op<string mnemonic, list<Trait> traits = []> :
+    Op<GEN_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// ND-Range Operations
+//===----------------------------------------------------------------------===//
+
+class GEN_3DNDRangeOp<string mnemonic, list<Trait> traits = []>
+    : GEN_Op<mnemonic, [NoMemoryEffect] # traits> {
+  let arguments = (ins I32:$dim);
+  let results = (outs Index:$res);
+
+  let assemblyFormat = "$dim attr-dict";
+}
+
+def GEN_LocalIdOp : GEN_3DNDRangeOp<"local_id"> {
+  let summary = "Query a work-item's local id.";
+  let description = [{
+    Query the work-item's position in its work-group, i.e., its local id, in a
+    given dimension `dim`, which must be either 0, 1 or 2. Behavior is undefined
+    for invalid `dim` values.
+
+    Example:
+    ```mlir
+    %local_id = gen.local_id %dim
+    ```
+  }];
+}
+
+def GEN_WorkGroupIdOp : GEN_3DNDRangeOp<"work_group_id"> {
+  let summary = "Query the id of a work-item's work-group.";
+  let description = [{
+    Query the id of a work-item's work-group in a given dimension `dim`, which
+    must be either 0, 1 or 2. Behavior is undefined for invalid `dim` values.
+
+    Example:
+    ```mlir
+    %work_group_id = gen.work_group_id %dim
+    ```
+  }];
+}
+
+def GEN_WorkGroupSizeOp : GEN_3DNDRangeOp<"work_group_size"> {
+  let summary = "Query the work-group size.";
+  let description = [{
+    Query the work-item's work-group size in a given dimension `dim`, which must
+    be either 0, 1 or 2. Behavior is undefined for invalid `dim` values.
+
+    Example:
+    ```mlir
+    %work_group_size = gen.work_group_size %dim
+    ```
+  }];
+}
+
+def GEN_NumWorkGroupsOp : GEN_3DNDRangeOp<"num_work_groups"> {
+  let summary = "Query the number of work-groups in the ND-range.";
+  let description = [{
+    Query the number of work-groups in the ND-range in a given dimension `dim`,
+    which must be either 0, 1 or 2. Behavior is undefined for invalid `dim`
+    values.
+
+    Example:
+    ```mlir
+    %wg_number = gen.num_work_groups %dim
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Synchronization
+//===----------------------------------------------------------------------===//
+
+def GEN_BarrierOp : GEN_Op<"barrier"> {
+  let summary = "Synchronizes all work-items of a work-group.";
+
+  let description = [{
+    Wait for all work-items in a given work-group to reach this execution point.
+    All memory accesses made by these work-items prior to the operation are
+    visible to all work-items in the work-group.
+
+    It is undefined behavior unless none or all work-items in the work-group
+    reach this execution point.
+
+    Example:
+    ```mlir
+    gen.barrier
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+def ShuffleValueType
+    : AnyTypeOf<[SignlessIntOfWidths<[8, 16, 32, 64]>, FloatOfWidths<[16, 32, 64]>]>;
+
+def GEN_SubGroupShuffleOp
+    : GEN_Op<"sub_group_shuffle", [Pure, AllTypesMatch<["res", "value"]>]> {
+  let summary = "Sub-group shuffle";
+  let description = [{
+    The `gen.sub_group_shuffle` operation is invoked by different work items
+    with different values, given by `value`. Different work items have different
+    sub-group local IDs. The shuffle kind, `kind`, is given to determine how to
+    calculate the associated sub-group local ID. It returns the associated
+    `value` for the work item with sub-group local ID equal to:
+    - `kind` == xor, the current invocation’s sub-group local ID xor'ed with `mask`.
+    - `kind` == up, the current invocation’s sub-group local ID - `mask`.
+    - `kind` == down, the current invocation’s sub-group local ID + `mask`.
+    - `kind` == idx, the sub-group local ID `mask`.
+
+    `value` and `res` types must match and can be any of: `i8`, `i16`, `i32`,
+    `i64`, `f16`, `f32` or `f64`.
+
+    Example:
+    ```mlir
+    // xor shuffle
+    %0 = gen.sub_group_shuffle xor %arg0, %arg4 : i32
+
+    // up shuffle
+    %1 = gen.sub_group_shuffle up %arg1, %arg4 : i64
+
+    // down shuffle
+    %2 = gen.sub_group_shuffle down %arg2, %arg4 : f32
+
+    // idx shuffle
+    %3 = gen.sub_group_shuffle idx %arg3, %arg4 : f64
+    ```
+  }];
+
+  let arguments = (ins ShuffleValueType:$value,
+                       I32:$mask,
+                       GEN_ShflKindAttr:$kind);
+  let results = (outs ShuffleValueType:$res);
+
+  let assemblyFormat = [{
+    $kind $value `,` $mask attr-dict `:` type($res)
+  }];
+}
+
+#endif // GEN_OPS

mlir/include/mlir/InitAllDialects.h

@@ -36,6 +36,7 @@
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GEN/IR/GENDialect.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h"
 #include "mlir/Dialect/IRDL/IR/IRDL.h"
@@ -116,6 +117,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   DLTIDialect,
                   emitc::EmitCDialect,
                   func::FuncDialect,
+                  GEN::GENDialect,
                   gpu::GPUDialect,
                   index::IndexDialect,
                   irdl::IRDLDialect,

mlir/lib/Dialect/CMakeLists.txt

@@ -12,6 +12,7 @@ add_subdirectory(ControlFlow)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(Func)
+add_subdirectory(GEN)
 add_subdirectory(GPU)
 add_subdirectory(Index)
 add_subdirectory(IRDL)

mlir/lib/Dialect/GEN/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,15 @@
+add_mlir_dialect_library(MLIRGENDialect
+  GENDialect.cpp
+  GENOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GEN
+
+  DEPENDS
+  MLIRGENOpsIncGen
+  MLIRGENOpsEnumsIncGen
+  MLIRGENOpsAttrDefsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+)

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

@@ -0,0 +1,34 @@
+//===- GENDialect.cpp - MLIR GEN Dialect implementation -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GEN/IR/GENDialect.h"
+
+#include "mlir/Dialect/GEN/IR/GENOps.h"
+#include "mlir/IR/DialectImplementation.h"
+
+using namespace mlir;
+using namespace mlir::GEN;
+
+#include "mlir/Dialect/GEN/IR/GENOpsDialect.cpp.inc"
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// GEN dialect.
+//===----------------------------------------------------------------------===//
+
+void GENDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/GEN/IR/GENOps.cpp.inc"
+      >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/GEN/IR/GENOpsAttrDefs.cpp.inc"
+      >();
+}

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

@@ -0,0 +1,24 @@
+//===- GENOps.cpp - GEN dialect operations --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GEN/IR/GENOps.h"
+
+#include "mlir/IR/Builders.h"
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/GEN/IR/GENOps.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// TableGen'd enum attribute definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GEN/IR/GENOpsEnums.cpp.inc"