[Mlir-commits] [mlir] [MLIR][GEN] Add GEN dialect (PR #88734)

[Jakub Kuderski]

================
@@ -0,0 +1,133 @@
+//===-- 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>;
+
+class GENOpTrait<string name, list<Trait> traits = [],
+                 code extraOpDeclaration = [{}],
+                 code extraOpDefinition = [{}]>
+    : NativeOpTrait<name, traits, extraOpDeclaration, extraOpDefinition> {
+  let cppNamespace = "::mlir::OpTrait::GEN";
+}
+
+//===----------------------------------------------------------------------===//
+// ND-Range Operations
+//===----------------------------------------------------------------------===//
+
+def GEN3DNDRange : GENOpTrait<"GEN3DNDRange">;
+
+class GEN_3DNDRangeOp<string mnemonic, list<Trait> traits = []>
+    : GEN_Op<mnemonic, [GEN3DNDRange, Pure] # traits>,
+      Arguments<(ins I32:$dim)>,
+      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.
+    ```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.
+    ```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.
+    ```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.
+    ```mlir
+    %wg_number = gen.num_work_groups %dim
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Synchronization
+//===----------------------------------------------------------------------===//
+
+def GEN_BarrierOp : GEN_Op<"barrier"> {
+  let summary = "Work-group barrier";
+
+  let description = [{
+    The `gen.barrier` operation performs a work-group barrier and ensures all
+    outstanding memory transaction using local or global memory are complete.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+def IntegerOrFloatType : AnyTypeOf<[AnySignlessInteger, AnyFloat]>;
+
+def GEN_SubGroupShuffleOp : GEN_Op<"sub_group_shuffle", [
+      TypesMatchWith<"result and value have the same type",
+                     "res", "value", "$_self">]>,
+  Results<(outs IntegerOrFloatType:$res)>,
+  Arguments<(ins IntegerOrFloatType:$value,
+                 I32:$mask,
+                 GEN_ShflKindAttr:$kind)> {
----------------
kuhar wrote:

nit: I find it more readable to define results and arguments in the op definition body, e.g.:
https://github.com/llvm/llvm-project/blob/a3ce29f7bb5510131971ac5ccc63132dd48c8dd2/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td#L29-L37

https://github.com/llvm/llvm-project/pull/88734