[SPIRV] Support for SPV_INTEL_ternary_bitwise_function (llvm#134866)

Adds support for the SPV_INTEL_ternary_bitwise_function extension,
adding;
* the OpBitwiseFunctionINTEL SPIR-V instruction, a ternary bitwise
function where the operation performed is determined by a look-up table
index,
 * and the corresponding TernaryBitwiseFunctionINTEL capability.

See
https://github.khronos.org/SPIRV-Registry/extensions/INTEL/SPV_INTEL_ternary_bitwise_function.html.

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

llvm/docs/SPIRVUsage.rst

@@ -209,6 +209,8 @@ list of supported SPIR-V extensions, sorted alphabetically by their extension na
      - Adds the ability to declare extended instruction sets that have no semantic impact and can be safely removed from a module.
    * - ``SPV_INTEL_fp_max_error``
      - Adds the ability to specify the maximum error for floating-point operations.
+   * - ``SPV_INTEL_ternary_bitwise_function``
+     - Adds a bitwise instruction on three operands and a look-up table index for specifying the bitwise operation to perform. 
 
 To enable multiple extensions, list them separated by comma. For example, to enable support for atomic operations on floating-point numbers and arbitrary precision integers, use:
 

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

@@ -1058,6 +1058,24 @@ static bool buildBindlessImageINTELInst(const SPIRV::IncomingCall *Call,
   return true;
 }
 
+/// Helper function for building Intel's OpBitwiseFunctionINTEL instruction.
+static bool buildTernaryBitwiseFunctionINTELInst(
+    const SPIRV::IncomingCall *Call, unsigned Opcode,
+    MachineIRBuilder &MIRBuilder, SPIRVGlobalRegistry *GR) {
+  // Generate SPIRV instruction accordingly.
+  if (Call->isSpirvOp())
+    return buildOpFromWrapper(MIRBuilder, Opcode, Call,
+                              GR->getSPIRVTypeID(Call->ReturnType));
+
+  auto MIB = MIRBuilder.buildInstr(Opcode)
+                 .addDef(Call->ReturnRegister)
+                 .addUse(GR->getSPIRVTypeID(Call->ReturnType));
+  for (unsigned i = 0; i < Call->Arguments.size(); ++i)
+    MIB.addUse(Call->Arguments[i]);
+
+  return true;
+}
+
 static unsigned getNumComponentsForDim(SPIRV::Dim::Dim dim) {
   switch (dim) {
   case SPIRV::Dim::DIM_1D:
@@ -2264,6 +2282,18 @@ static bool generateBindlessImageINTELInst(const SPIRV::IncomingCall *Call,
   return buildBindlessImageINTELInst(Call, Opcode, MIRBuilder, GR);
 }
 
+static bool
+generateTernaryBitwiseFunctionINTELInst(const SPIRV::IncomingCall *Call,
+                                        MachineIRBuilder &MIRBuilder,
+                                        SPIRVGlobalRegistry *GR) {
+  // Lookup the instruction opcode in the TableGen records.
+  const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
+  unsigned Opcode =
+      SPIRV::lookupNativeBuiltin(Builtin->Name, Builtin->Set)->Opcode;
+
+  return buildTernaryBitwiseFunctionINTELInst(Call, Opcode, MIRBuilder, GR);
+}
+
 static bool buildNDRange(const SPIRV::IncomingCall *Call,
                          MachineIRBuilder &MIRBuilder,
                          SPIRVGlobalRegistry *GR) {
@@ -2845,6 +2875,8 @@ std::optional<bool> lowerBuiltin(const StringRef DemangledCall,
     return generateExtendedBitOpsInst(Call.get(), MIRBuilder, GR);
   case SPIRV::BindlessINTEL:
     return generateBindlessImageINTELInst(Call.get(), MIRBuilder, GR);
+  case SPIRV::TernaryBitwiseINTEL:
+    return generateTernaryBitwiseFunctionINTELInst(Call.get(), MIRBuilder, GR);
   }
   return false;
 }

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

@@ -67,6 +67,7 @@ def CoopMatr : BuiltinGroup;
 def ICarryBorrow : BuiltinGroup;
 def ExtendedBitOps : BuiltinGroup;
 def BindlessINTEL : BuiltinGroup;
+def TernaryBitwiseINTEL : BuiltinGroup;
 
 //===----------------------------------------------------------------------===//
 // Class defining a demangled builtin record. The information in the record
@@ -714,6 +715,9 @@ defm : DemangledNativeBuiltin<"__spirv_ConvertHandleToImageINTEL", OpenCL_std, B
 defm : DemangledNativeBuiltin<"__spirv_ConvertHandleToSamplerINTEL", OpenCL_std, BindlessINTEL, 1, 1, OpConvertHandleToSamplerINTEL>;
 defm : DemangledNativeBuiltin<"__spirv_ConvertHandleToSampledImageINTEL", OpenCL_std, BindlessINTEL, 1, 1, OpConvertHandleToSampledImageINTEL>;
 
+// SPV_INTEL_ternary_bitwise_function builtin records:
+defm : DemangledNativeBuiltin<"__spirv_BitwiseFunctionINTEL", OpenCL_std, TernaryBitwiseINTEL, 4, 4, OpBitwiseFunctionINTEL>;
+
 //===----------------------------------------------------------------------===//
 // Class defining a work/sub group builtin that should be translated into a
 // SPIR-V instruction using the defined properties.

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -92,7 +92,9 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
         {"SPV_INTEL_long_composites",
          SPIRV::Extension::Extension::SPV_INTEL_long_composites},
         {"SPV_INTEL_fp_max_error",
-         SPIRV::Extension::Extension::SPV_INTEL_fp_max_error}};
+         SPIRV::Extension::Extension::SPV_INTEL_fp_max_error},
+        {"SPV_INTEL_ternary_bitwise_function",
+         SPIRV::Extension::Extension::SPV_INTEL_ternary_bitwise_function}};
 
 bool SPIRVExtensionsParser::parse(cl::Option &O, StringRef ArgName,
                                   StringRef ArgValue,

llvm/lib/Target/SPIRV/SPIRVInstrInfo.td

@@ -928,3 +928,7 @@ def OpAliasScopeDeclINTEL: Op<5912, (outs ID:$res), (ins ID:$AliasDomain, variab
                   "$res = OpAliasScopeDeclINTEL $AliasDomain">;
 def OpAliasScopeListDeclINTEL: Op<5913, (outs ID:$res), (ins variable_ops),
                   "$res = OpAliasScopeListDeclINTEL">;
+
+// SPV_INTEL_ternary_bitwise_function
+def OpBitwiseFunctionINTEL: Op<6242, (outs ID:$res), (ins TYPE:$type, ID:$a, ID:$b, ID:$c, ID:$lut_index),
+                  "$res = OpBitwiseFunctionINTEL $type $a $b $c $lut_index">;

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -1799,6 +1799,17 @@ void addInstrRequirements(const MachineInstr &MI,
     Reqs.addCapability(SPIRV::Capability::LongCompositesINTEL);
     break;
   }
+  case SPIRV::OpBitwiseFunctionINTEL: {
+    if (!ST.canUseExtension(
+            SPIRV::Extension::SPV_INTEL_ternary_bitwise_function))
+      report_fatal_error(
+          "OpBitwiseFunctionINTEL instruction requires the following SPIR-V "
+          "extension: SPV_INTEL_ternary_bitwise_function",
+          false);
+    Reqs.addExtension(SPIRV::Extension::SPV_INTEL_ternary_bitwise_function);
+    Reqs.addCapability(SPIRV::Capability::TernaryBitwiseFunctionINTEL);
+    break;
+  }
 
   default:
     break;

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -313,6 +313,7 @@ defm SPV_INTEL_bindless_images : ExtensionOperand<116>;
 defm SPV_INTEL_long_composites : ExtensionOperand<117>;
 defm SPV_INTEL_memory_access_aliasing : ExtensionOperand<118>;
 defm SPV_INTEL_fp_max_error : ExtensionOperand<119>;
+defm SPV_INTEL_ternary_bitwise_function : ExtensionOperand<120>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
@@ -513,6 +514,7 @@ defm LongCompositesINTEL : CapabilityOperand<6089, 0, 0, [SPV_INTEL_long_composi
 defm BindlessImagesINTEL : CapabilityOperand<6528, 0, 0, [SPV_INTEL_bindless_images], []>;
 defm MemoryAccessAliasingINTEL : CapabilityOperand<5910, 0, 0, [SPV_INTEL_memory_access_aliasing], []>;
 defm FPMaxErrorINTEL : CapabilityOperand<6169, 0, 0, [SPV_INTEL_fp_max_error], []>;
+defm TernaryBitwiseFunctionINTEL : CapabilityOperand<6241, 0, 0, [SPV_INTEL_ternary_bitwise_function], []>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define SourceLanguage enum values and at the same time

llvm/test/CodeGen/SPIRV/extensions/SPV_INTEL_ternary_bitwise_function/bitwise-function.ll

@@ -0,0 +1,58 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv32-unknown-unknown %s --spirv-ext=+SPV_INTEL_ternary_bitwise_function -o - | FileCheck %s --check-prefix=CHECK-EXTENSION
+; RUN: not llc -verify-machineinstrs -O0 -mtriple=spirv32-unknown-unknown %s -o %t.spvt 2>&1 | FileCheck %s --check-prefix=CHECK-NO-EXTENSION
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s --spirv-ext=+SPV_INTEL_ternary_bitwise_function -o - -filetype=obj | spirv-val %} 
+;
+; CHECK-NO-EXTENSION: LLVM ERROR: OpBitwiseFunctionINTEL instruction requires the following SPIR-V extension: SPV_INTEL_ternary_bitwise_function
+;
+; CHECK-EXTENSION-NOT: Name [[#]] "_Z28__spirv_BitwiseFunctionINTELiiij"
+; CHECK-EXTENSION-NOT: Name [[#]] "_Z28__spirv_BitwiseFunctionINTELDv4_iS_S_j"
+;
+; CHECK-EXTENSION-DAG: Capability TernaryBitwiseFunctionINTEL
+; CHECK-EXTENSION-DAG: Extension "SPV_INTEL_ternary_bitwise_function"
+; CHECK-EXTENSION-DAG: %[[#TYPEINT:]] = OpTypeInt 32 0
+; CHECK-EXTENSION-DAG: %[[#TYPEINTVEC4:]] = OpTypeVector %[[#TYPEINT]] 4
+; CHECK-EXTENSION-DAG: %[[#ScalarLUT:]] = OpConstant %[[#TYPEINT]] 24
+; CHECK-EXTENSION-DAG: %[[#VecLUT:]] = OpConstant %[[#TYPEINT]] 42
+; CHECK-EXTENSION: %[[#ScalarA:]] = OpLoad %[[#TYPEINT]]
+; CHECK-EXTENSION: %[[#ScalarB:]] = OpLoad %[[#TYPEINT]]
+; CHECK-EXTENSION: %[[#ScalarC:]] = OpLoad %[[#TYPEINT]]
+; CHECK-EXTENSION: %{{.*}}  = OpBitwiseFunctionINTEL %[[#TYPEINT]] %[[#ScalarA]] %[[#ScalarB]] %[[#ScalarC]] %[[#ScalarLUT]]
+; CHECK-EXTENSION: %[[#VecA:]] = OpLoad %[[#TYPEINTVEC4]]
+; CHECK-EXTENSION: %[[#VecB:]] = OpLoad %[[#TYPEINTVEC4]]
+; CHECK-EXTENSION: %[[#VecC:]] = OpLoad %[[#TYPEINTVEC4]]
+; CHECK-EXTENSION: %{{.*}}  = OpBitwiseFunctionINTEL %[[#TYPEINTVEC4]] %[[#VecA]] %[[#VecB]] %[[#VecC]] %[[#VecLUT]]
+
+; Function Attrs: nounwind readnone
+define spir_kernel void @fooScalar() {
+entry:
+  %argA = alloca i32
+  %argB = alloca i32
+  %argC = alloca i32
+  %A = load i32, ptr %argA
+  %B = load i32, ptr %argB
+  %C = load i32, ptr %argC
+  %res = call spir_func i32 @_Z28__spirv_BitwiseFunctionINTELiiii(i32 %A, i32 %B, i32 %C, i32 24)
+  ret void
+}
+
+; Function Attrs: nounwind readnone
+define spir_kernel void @fooVec() {
+entry:
+  %argA = alloca <4 x i32>
+  %argB = alloca <4 x i32>
+  %argC = alloca <4 x i32>
+  %A = load <4 x i32>, ptr %argA
+  %B = load <4 x i32>, ptr %argB
+  %C = load <4 x i32>, ptr %argC
+  %res = call spir_func <4 x i32> @_Z28__spirv_BitwiseFunctionINTELDv4_iS_S_i(<4 x i32> %A, <4 x i32> %B, <4 x i32> %C, i32 42)
+  ret void
+}
+
+declare dso_local spir_func i32 @_Z28__spirv_BitwiseFunctionINTELiiii(i32, i32, i32, i32)
+declare dso_local spir_func <4 x i32> @_Z28__spirv_BitwiseFunctionINTELDv4_iS_S_i(<4 x i32>, <4 x i32>, <4 x i32>, i32)
+
+!llvm.module.flags = !{!0}
+!opencl.spir.version = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 1, i32 2}