[mlir][spirv] Add definition for OpKill (#126554)

Although the operation is deprecated in the most recent version of the
SPIR-V spec, it is still used by older shaders, so having it defined is
valuable and incurs negligible maintenance overhead, due to op
simplicity.

Author: IgWod-IMG

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

@@ -4445,6 +4445,7 @@ def SPIRV_OC_OpSelectionMerge                : I32EnumAttrCase<"OpSelectionMerge
 def SPIRV_OC_OpLabel                         : I32EnumAttrCase<"OpLabel", 248>;
 def SPIRV_OC_OpBranch                        : I32EnumAttrCase<"OpBranch", 249>;
 def SPIRV_OC_OpBranchConditional             : I32EnumAttrCase<"OpBranchConditional", 250>;
+def SPIRV_OC_OpKill                          : I32EnumAttrCase<"OpKill", 252>;
 def SPIRV_OC_OpReturn                        : I32EnumAttrCase<"OpReturn", 253>;
 def SPIRV_OC_OpReturnValue                   : I32EnumAttrCase<"OpReturnValue", 254>;
 def SPIRV_OC_OpUnreachable                   : I32EnumAttrCase<"OpUnreachable", 255>;
@@ -4574,7 +4575,7 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpAtomicAnd, SPIRV_OC_OpAtomicOr, SPIRV_OC_OpAtomicXor,
       SPIRV_OC_OpPhi, SPIRV_OC_OpLoopMerge, SPIRV_OC_OpSelectionMerge,
       SPIRV_OC_OpLabel, SPIRV_OC_OpBranch, SPIRV_OC_OpBranchConditional,
-      SPIRV_OC_OpReturn, SPIRV_OC_OpReturnValue, SPIRV_OC_OpUnreachable,
+      SPIRV_OC_OpKill, SPIRV_OC_OpReturn, SPIRV_OC_OpReturnValue, SPIRV_OC_OpUnreachable,
       SPIRV_OC_OpGroupBroadcast, SPIRV_OC_OpGroupIAdd, SPIRV_OC_OpGroupFAdd,
       SPIRV_OC_OpGroupFMin, SPIRV_OC_OpGroupUMin, SPIRV_OC_OpGroupSMin,
       SPIRV_OC_OpGroupFMax, SPIRV_OC_OpGroupUMax, SPIRV_OC_OpGroupSMax,

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td

@@ -242,6 +242,48 @@ def SPIRV_FunctionCallOp : SPIRV_Op<"FunctionCall", [
 
 // -----
 
+def SPIRV_KillOp : SPIRV_Op<"Kill", [Terminator]> {
+  let summary = [{
+    Deprecated (use OpTerminateInvocation or OpDemoteToHelperInvocation).
+  }];
+
+  let description = [{
+    Fragment-shader discard.
+
+    Ceases all further processing in any invocation that executes it: Only
+    instructions these invocations executed before OpKill have observable
+    side effects. If this instruction is executed in non-uniform control
+    flow, all subsequent control flow is non-uniform (for invocations that
+    continue to execute).
+
+    This instruction must be the last instruction in a block.
+
+    This instruction is only valid in the Fragment Execution Model.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    spirv.Kill
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[]>,
+    Capability<[SPIRV_C_Shader]>
+  ];
+
+  let arguments = (ins);
+  let results = (outs);
+  let assemblyFormat = "attr-dict";
+  let hasVerifier = 0;
+}
+
+// -----
+
 def SPIRV_LoopOp : SPIRV_Op<"mlir.loop", [InFunctionScope]> {
   let summary = "Define a structured loop.";
 

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

@@ -84,7 +84,11 @@ struct SPIRVInlinerInterface : public DialectInlinerInterface {
     // TODO: we need to filter OpKill here to avoid inlining it to
     // a loop continue construct:
     // https://github.com/KhronosGroup/SPIRV-Headers/issues/86
-    // However OpKill is fragment shader specific and we don't support it yet.
+    // For now, we just disallow inlining OpKill anywhere in the code,
+    // but this restriction should be relaxed, as pointed above.
+    if (isa<spirv::KillOp>(op))
+      return false;
+
     return true;
   }
 

mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir

@@ -789,3 +789,15 @@ func.func @unreachable() {
   // expected-error @+1 {{cannot be used in reachable block}}
   spirv.Unreachable
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.Kill
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @kill
+func.func @kill() {
+  // CHECK: spirv.Kill
+  spirv.Kill
+}

mlir/test/Dialect/SPIRV/Transforms/inlining.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(spirv.module(inline{default-pipeline=''}))' | FileCheck %s
+// RUN: mlir-opt %s --split-input-file --pass-pipeline='builtin.module(spirv.module(inline{default-pipeline=''}))' | FileCheck %s
 
 spirv.module Logical GLSL450 {
   spirv.func @callee() "None" {
@@ -246,5 +246,24 @@ spirv.module Logical GLSL450 {
   }
 }
 
+// -----
+
+spirv.module Logical GLSL450 {
+  // CHECK-LABEL: @callee
+  spirv.func @callee() -> () "None" {
+    // CHECK-NEXT: spirv.Kill
+    spirv.Kill
+  }
+
+  // CHECK-LABEL: @do_not_inline_kill
+  spirv.func @do_not_inline_kill() -> () "None" {
+    // CHECK-NOT: spirv.Kill
+    // CHECK-NEXT: spirv.FunctionCall @callee() : () -> ()
+    spirv.FunctionCall @callee() : () -> ()
+    // CHECK-NEXT: spirv.Return
+    spirv.Return
+  }
+}
+
 // TODO: Add tests for inlining structured control flow into
 // structured control flow.

mlir/test/Target/SPIRV/terminator.mlir

@@ -24,4 +24,10 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
     // CHECK-NOT: spirv.Unreachable
     spirv.Unreachable
   }
+
+  // CHECK-LABEL: @kill
+  spirv.func @kill() -> () "None" {
+    // CHECK: spirv.Kill
+    spirv.Kill
+  }
 }