[mlir][spirv] Add definition for OpEmitVertex and OpEndPrimitive (#123759)

This is hopefully the first patch in the series of patches adding some
missing SPIR-V ops to MLIR over the next weeks/months, starting with
something simple: `OpEmitVertex` and `OpEndPrimitive`. Since the ops
have no input and outputs, and the only condition is "This instruction
must only be used when only one stream is present.", which I don't think
can be validate at the instruction level in isolation, I set
`hasVerifier` to 0. I hope I didn't miss anything, but I'm more than
happy to address any comments.

6 files changed, 148 insertions, 4 deletions

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index 469a9a0..c84677d 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4438,6 +4438,8 @@ def SPIRV_OC_OpBitFieldSExtract             : I32EnumAttrCase<"OpBitFieldSExtrac
 def SPIRV_OC_OpBitFieldUExtract             : I32EnumAttrCase<"OpBitFieldUExtract", 203>;
 def SPIRV_OC_OpBitReverse                   : I32EnumAttrCase<"OpBitReverse", 204>;
 def SPIRV_OC_OpBitCount                     : I32EnumAttrCase<"OpBitCount", 205>;
+def SPIRV_OC_OpEmitVertex                   : I32EnumAttrCase<"OpEmitVertex", 218>;
+def SPIRV_OC_OpEndPrimitive                 : I32EnumAttrCase<"OpEndPrimitive", 219>;
 def SPIRV_OC_OpControlBarrier               : I32EnumAttrCase<"OpControlBarrier", 224>;
 def SPIRV_OC_OpMemoryBarrier                : I32EnumAttrCase<"OpMemoryBarrier", 225>;
 def SPIRV_OC_OpAtomicExchange               : I32EnumAttrCase<"OpAtomicExchange", 229>;
@@ -4576,7 +4578,8 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpBitwiseOr, SPIRV_OC_OpBitwiseXor, SPIRV_OC_OpBitwiseAnd,
       SPIRV_OC_OpNot, SPIRV_OC_OpBitFieldInsert, SPIRV_OC_OpBitFieldSExtract,
       SPIRV_OC_OpBitFieldUExtract, SPIRV_OC_OpBitReverse, SPIRV_OC_OpBitCount,
-      SPIRV_OC_OpControlBarrier, SPIRV_OC_OpMemoryBarrier, SPIRV_OC_OpAtomicExchange,
+      SPIRV_OC_OpEmitVertex, SPIRV_OC_OpEndPrimitive, SPIRV_OC_OpControlBarrier,
+      SPIRV_OC_OpMemoryBarrier, SPIRV_OC_OpAtomicExchange,
       SPIRV_OC_OpAtomicCompareExchange, SPIRV_OC_OpAtomicCompareExchangeWeak,
       SPIRV_OC_OpAtomicIIncrement, SPIRV_OC_OpAtomicIDecrement,
       SPIRV_OC_OpAtomicIAdd, SPIRV_OC_OpAtomicISub, SPIRV_OC_OpAtomicSMin,
@@ -4609,9 +4612,8 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpCooperativeMatrixLengthKHR, SPIRV_OC_OpSubgroupBlockReadINTEL,
       SPIRV_OC_OpSubgroupBlockWriteINTEL, SPIRV_OC_OpAssumeTrueKHR,
       SPIRV_OC_OpAtomicFAddEXT, SPIRV_OC_OpConvertFToBF16INTEL,
-      SPIRV_OC_OpConvertBF16ToFINTEL,
-      SPIRV_OC_OpControlBarrierArriveINTEL, SPIRV_OC_OpControlBarrierWaitINTEL,
-      SPIRV_OC_OpGroupIMulKHR,
+      SPIRV_OC_OpConvertBF16ToFINTEL, SPIRV_OC_OpControlBarrierArriveINTEL,
+      SPIRV_OC_OpControlBarrierWaitINTEL, SPIRV_OC_OpGroupIMulKHR,
       SPIRV_OC_OpGroupFMulKHR
     ]>;
 
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
index 9912f195..ff1ca89 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
@@ -40,6 +40,7 @@ include "mlir/Dialect/SPIRV/IR/SPIRVMatrixOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVPrimitiveOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVCLOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVPrimitiveOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVPrimitiveOps.td
new file mode 100755
index 0000000..c390ae5
--- /dev/null
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVPrimitiveOps.td
@@ -0,0 +1,81 @@
+//===-- SPIRVPrimitiveOps.td - MLIR SPIR-V Primitive Ops ------*- 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 file contains primitive ops for the SPIR-V dialect. It corresponds
+// to "3.52.19. Primitive Instructions" of the SPIR-V specification.
+//
+//===-----------------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_PRIMITIVE_OPS
+#define MLIR_DIALECT_SPIRV_PRIMITIVE_OPS
+
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+
+// -----
+
+def SPIRV_EmitVertexOp : SPIRV_Op<"EmitVertex", []> {
+  let summary = [{
+    Emits the current values of all output variables to the current output
+    primitive. After execution, the values of all output variables are
+    undefined.
+  }];
+
+  let description = [{
+    This instruction must only be used when only one stream is present.
+
+    #### Example:
+
+    ```mlir
+    spirv.EmitVertex
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[]>,
+    Capability<[SPIRV_C_Geometry]>
+  ];
+
+  let arguments = (ins);
+  let results = (outs);
+  let hasVerifier = 0;
+  let assemblyFormat = "attr-dict";
+}
+
+// -----
+
+def SPIRV_EndPrimitiveOp : SPIRV_Op<"EndPrimitive", []> {
+  let summary = [{
+    Finish the current primitive and start a new one. No vertex is emitted.
+  }];
+
+  let description = [{
+    This instruction must only be used when only one stream is present.
+
+    #### Example:
+
+    ```mlir
+    spirv.EndPrimitive
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[]>,
+    Capability<[SPIRV_C_Geometry]>
+  ];
+
+  let arguments = (ins);
+  let results = (outs);
+  let hasVerifier = 0;
+  let assemblyFormat = "attr-dict";
+}
+
+#endif // MLIR_DIALECT_SPIRV_PRIMITIVE_OPS
diff --git a/mlir/test/Dialect/SPIRV/IR/availability.mlir b/mlir/test/Dialect/SPIRV/IR/availability.mlir
index c583a48..31a90ad 100644
--- a/mlir/test/Dialect/SPIRV/IR/availability.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/availability.mlir
@@ -233,3 +233,25 @@ func.func @udot_acc_sat_vector_4xi16_i64(%a: vector<4xi16>, %acc: i64) -> i64 {
   %r = spirv.UDotAccSat %a, %a, %acc: vector<4xi16> -> i64
   return %r: i64
 }
+
+//===----------------------------------------------------------------------===//
+// Primitive ops
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: emit_vertex
+func.func @emit_vertex() -> () {
+  // CHECK: min version: v1.0
+  // CHECK: max version: v1.6
+  // CHECK: capabilities: [ [Geometry] ]
+  spirv.EmitVertex
+  return
+}
+
+// CHECK-LABEL: end_primitive
+func.func @end_primitive() -> () {
+  // CHECK: min version: v1.0
+  // CHECK: max version: v1.6
+  // CHECK: capabilities: [ [Geometry] ]
+  spirv.EndPrimitive
+  return
+}
diff --git a/mlir/test/Dialect/SPIRV/IR/primitive-ops.mlir b/mlir/test/Dialect/SPIRV/IR/primitive-ops.mlir
new file mode 100644
index 0000000..451c334
--- /dev/null
+++ b/mlir/test/Dialect/SPIRV/IR/primitive-ops.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spirv.EmitVertex
+//===----------------------------------------------------------------------===//
+
+func.func @emit_vertex() {
+  // CHECK: spirv.EmitVertex
+  spirv.EmitVertex
+  spirv.Return
+}
+
+//===----------------------------------------------------------------------===//
+// spirv.EndPrimitive
+//===----------------------------------------------------------------------===//
+
+func.func @end_primitive() {
+  // CHECK: spirv.EndPrimitive
+  spirv.EndPrimitive
+  spirv.Return
+}
diff --git a/mlir/test/Target/SPIRV/primitive-ops.mlir b/mlir/test/Target/SPIRV/primitive-ops.mlir
new file mode 100644
index 0000000..63a0b1e
--- /dev/null
+++ b/mlir/test/Target/SPIRV/primitive-ops.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip %s | FileCheck %s
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Geometry], []> {
+  spirv.GlobalVariable @out : !spirv.ptr<!spirv.struct<(vector<4xf32>, f32, !spirv.array<1 x f32>)>, Output> 
+  spirv.func @primitive_ops() "None" {
+    // CHECK: spirv.EmitVertex
+    spirv.EmitVertex
+    // CHECK: spirv.EndPrimitive
+    spirv.EndPrimitive
+    spirv.Return
+  }
+  spirv.EntryPoint "Geometry" @primitive_ops, @out
+  spirv.ExecutionMode @primitive_ops "InputPoints"
+  spirv.ExecutionMode @primitive_ops "Invocations", 1
+  spirv.ExecutionMode @primitive_ops "OutputLineStrip"
+  spirv.ExecutionMode @primitive_ops "OutputVertices", 2
+}