[SPIRV] Implement support for SPV_KHR_expect_assume

[pmatos]

Summary

Adds new extension SPV_KHR_expect_assume, new capability
ExpectAssumeKHR as well as the new instructions:

• OpExpectKHR
• OpAssumeTrueKHR

These are lowered from respectively llvm.expect.<ty> and llvm.assume
intrinsics.

[pmatos]

llvm/include/llvm/IR/IntrinsicsSPIRV.td

40

I will make sure to add a new line in the next revision of the patch.

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

57

I have a question here... We have these instruction groups. I added these to builtins but I feel it's not the right thing to do. I get a warning during compilation because I didn't handle these instructions:

/home/pmatos/dev/llvm-project/llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp:1933:11: warning: enumeration value 'ExpectAssume' not handled in switch [-Wswitch]
  switch (Call->Builtin->Group) {
          ^~~~~~~~~~~~~~~~~~~~
1 warning generated.

Should I create a group just for these instructions instead?

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

860

Other cases don't addExtensions here, just capabilities but if I don't then there's really no good place to add extensions. Would it make sense to have a loop to add extensions based on capabilities required? Or is it possible that two different extensions provide the same capability?

llvm/test/CodeGen/SPIRV/assume.ll

2

I added testing on spirv64, I assume this makes sense.

16

I will make sure to add a new line in the next revision of the patch.

[iliya-diyachkov]

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

57

I'm not sure where this new ExpectAssume is actually used. If you are going to use ExpectAssume builtin in later version of the patch (or in the next patch), please add it later.

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

860

The place looks good but we need to check ST.canUseExtension().

Or is it possible that two different extensions provide the same capability?

I'm not sure but Reqs should work as a set and do not duplicate items.

[pmatos]

Currently https://github.com/llvm/llvm-project/pull/66217

llvm/include/llvm/IR/IntrinsicsSPIRV.td

let TargetPrefix = "spv" in {
   def int_spv_insertelt : Intrinsic<[llvm_any_ty], [llvm_any_ty, llvm_any_ty, llvm_anyint_ty]>;
   def int_spv_const_composite : Intrinsic<[llvm_i32_ty], [llvm_vararg_ty]>;
   def int_spv_bitcast : Intrinsic<[llvm_any_ty], [llvm_any_ty]>;
   def int_spv_switch : Intrinsic<[], [llvm_any_ty, llvm_vararg_ty]>;
   def int_spv_cmpxchg : Intrinsic<[llvm_i32_ty], [llvm_any_ty, llvm_vararg_ty]>;
   def int_spv_unreachable : Intrinsic<[], []>;
   def int_spv_alloca : Intrinsic<[llvm_any_ty], []>;
   def int_spv_undef : Intrinsic<[llvm_i32_ty], []>;
-}
+
+  // Expect, Assume Intrinsics
+  def int_spv_assume : Intrinsic<[], [llvm_i1_ty]>;
+  def int_spv_expect : Intrinsic<[llvm_anyint_ty], [LLVMMatchType<0>, LLVMMatchType<0>]>;
+}
 No newline at end of file

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

 def WriteImage : BuiltinGroup;
 def SampleImage : BuiltinGroup;
 def Select : BuiltinGroup;
 def SpecConstant : BuiltinGroup;
 def Enqueue : BuiltinGroup;
 def AsyncCopy : BuiltinGroup;
 def VectorLoadStore : BuiltinGroup;
 def LoadStore : BuiltinGroup;
+def ExpectAssume : BuiltinGroup;
 
 //===----------------------------------------------------------------------===//
 // Class defining a demangled builtin record. The information in the record
 // should be used to expand the builtin into either native SPIR-V instructions
 // or an external call (in case of builtins without a direct mapping).
 //
 // name is the demangled name of the given builtin.
 // set specifies which external instruction set the builtin belongs to.

llvm/lib/Target/SPIRV/SPIRVInstrInfo.td

 def OpGroupNonUniformUMax: OpGroupNUGroup<"UMax", 357>;
 def OpGroupNonUniformFMax: OpGroupNUGroup<"FMax", 358>;
 def OpGroupNonUniformBitwiseAnd: OpGroupNUGroup<"BitwiseAnd", 359>;
 def OpGroupNonUniformBitwiseOr: OpGroupNUGroup<"BitwiseOr", 360>;
 def OpGroupNonUniformBitwiseXor: OpGroupNUGroup<"BitwiseXor", 361>;
 def OpGroupNonUniformLogicalAnd: OpGroupNUGroup<"LogicalAnd", 362>;
 def OpGroupNonUniformLogicalOr: OpGroupNUGroup<"LogicalOr", 363>;
 def OpGroupNonUniformLogicalXor: OpGroupNUGroup<"LogicalXor", 364>;
+
+// SPV_KHR_expect_assume : Expect assume instructions
+def OpAssumeTrueKHR: Op<5630, (outs), (ins ID:$cond), "OpAssumeTrueKHR $cond">;
+def OpExpectKHR: Op<5631, (outs ID:$res), (ins TYPE:$ty, ID:$val, ID:$expected), "$res = OpExpectKHR $ty $val $expected">;

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

 //===- SPIRVInstructionSelector.cpp ------------------------------*- 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 implements the targeting of the InstructionSelector class for
 // SPIRV.
 // TODO: This should be generated by TableGen.
 //
 //===----------------------------------------------------------------------===//
 
+#include "MCTargetDesc/SPIRVMCTargetDesc.h"
 #include "SPIRV.h"
 #include "SPIRVGlobalRegistry.h"
 #include "SPIRVInstrInfo.h"
 #include "SPIRVRegisterBankInfo.h"
 #include "SPIRVRegisterInfo.h"
 #include "SPIRVTargetMachine.h"
 #include "SPIRVUtils.h"
 #include "llvm/ADT/APFloat.h"
bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   }
   case Intrinsic::spv_cmpxchg:
     return selectAtomicCmpXchg(ResVReg, ResType, I);
   case Intrinsic::spv_unreachable:
     BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpUnreachable));
     break;
   case Intrinsic::spv_alloca:
     return selectFrameIndex(ResVReg, ResType, I);
+  case Intrinsic::spv_assume:
+    BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpAssumeTrueKHR))
+        .addUse(I.getOperand(1).getReg());
+    break;
+  case Intrinsic::spv_expect:
+    BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpExpectKHR))
+        .addDef(ResVReg)
+        .addUse(GR.getSPIRVTypeID(ResType))
+        .addUse(I.getOperand(2).getReg())
+        .addUse(I.getOperand(3).getReg());
+    break;
   default:
     llvm_unreachable("Intrinsic selection not implemented");
   }
   return true;
 }
 
 bool SPIRVInstructionSelector::selectFrameIndex(Register ResVReg,
                                                 const SPIRVType *ResType,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

 // and performs the global register numbering.
 //
 // The results of this analysis are used in AsmPrinter to rename registers
 // globally and to output required instructions at the module level.
 //
 //===----------------------------------------------------------------------===//
 
 #include "SPIRVModuleAnalysis.h"
+#include "MCTargetDesc/SPIRVBaseInfo.h"
+#include "MCTargetDesc/SPIRVMCTargetDesc.h"
 #include "SPIRV.h"
 #include "SPIRVSubtarget.h"
 #include "SPIRVTargetMachine.h"
 #include "SPIRVUtils.h"
 #include "TargetInfo/SPIRVTargetInfo.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
addAvailableCaps({Capability::DenormPreserve, Capability::DenormFlushToZero,
                       Capability::RoundingModeRTE,
                       Capability::RoundingModeRTZ});
   // TODO: verify if this needs some checks.
   addAvailableCaps({Capability::Float16, Capability::Float64});
 
   // Add cap for SPV_INTEL_optnone.
   // FIXME: this should be added only if the target has the extension.
   addAvailableCaps({Capability::OptNoneINTEL});
+  addAvailableCaps({Capability::ExpectAssumeKHR});
 
   // TODO: add OpenCL extensions.
 }
 } // namespace SPIRV
 } // namespace llvm
 
 // Add the required capabilities from a decoration instruction (including
 // BuiltIns).
void addInstrRequirements(const MachineInstr &MI,
   case SPIRV::OpGroupNonUniformBallot:
   case SPIRV::OpGroupNonUniformInverseBallot:
   case SPIRV::OpGroupNonUniformBallotBitExtract:
   case SPIRV::OpGroupNonUniformBallotBitCount:
   case SPIRV::OpGroupNonUniformBallotFindLSB:
   case SPIRV::OpGroupNonUniformBallotFindMSB:
     Reqs.addCapability(SPIRV::Capability::GroupNonUniformBallot);
     break;
+  case SPIRV::OpAssumeTrueKHR:
+  case SPIRV::OpExpectKHR:
+    Reqs.addExtension(SPIRV::Extension::SPV_KHR_expect_assume);
+    Reqs.addCapability(SPIRV::Capability::ExpectAssumeKHR);
+    break;
   default:
     break;
   }
 }
 
 static void collectReqs(const Module &M, SPIRV::ModuleAnalysisInfo &MAI,
                         MachineModuleInfo *MMI, const SPIRVSubtarget &ST) {
   // Collect requirements for existing instructions.

llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp

 //===----------------------------------------------------------------------===//
 
 #include "SPIRV.h"
 #include "SPIRVTargetMachine.h"
 #include "SPIRVUtils.h"
 #include "llvm/CodeGen/IntrinsicLowering.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsSPIRV.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
 
 using namespace llvm;
 
 namespace llvm {
 void initializeSPIRVPrepareFunctionsPass(PassRegistry &);
 }
static void buildUMulWithOverflowFunc(Function *UMulFunc) {
   // umul.with.overflow intrinsic return a structure, where the first element
   // is the multiplication result, and the second is an overflow bit.
   Type *StructTy = UMulFunc->getReturnType();
   Value *Agg = IRB.CreateInsertValue(PoisonValue::get(StructTy), Mul, {0});
   Value *Res = IRB.CreateInsertValue(Agg, Overflow, {1});
   IRB.CreateRet(Res);
 }
 
+static void lowerExpectAssume(IntrinsicInst *II) {
+  // If we cannot use the SPV_KHR_expect_assume extension, then we need to
+  // ignore the intrinsic and move on. It should be removed later on by LLVM.
+  // Otherwise we should lower the intrinsic to the corresponding SPIR-V
+  // instruction.
+  // For @llvm.assume we have OpAssumeTrueKHR.
+  // For @llvm.expect we have OpExpectKHR.
+  //
+  // We need to lower this into a builtin and then the builtin into a SPIR-V
+  // instruction.
+  if (II->getIntrinsicID() == Intrinsic::assume) {
+    Function *F = Intrinsic::getDeclaration(
+        II->getModule(), Intrinsic::SPVIntrinsics::spv_assume);
+    II->setCalledFunction(F);
+  } else if (II->getIntrinsicID() == Intrinsic::expect) {
+    Function *F = Intrinsic::getDeclaration(
+        II->getModule(), Intrinsic::SPVIntrinsics::spv_expect,
+        {II->getOperand(0)->getType()});
+    II->setCalledFunction(F);
+  } else {
+    llvm_unreachable("Unknown intrinsic");
+  }
+
+  return;
+}
+
 static void lowerUMulWithOverflow(IntrinsicInst *UMulIntrinsic) {
   // Get a separate function - otherwise, we'd have to rework the CFG of the
   // current one. Then simply replace the intrinsic uses with a call to the new
   // function.
   Module *M = UMulIntrinsic->getModule();
   FunctionType *UMulFuncTy = UMulIntrinsic->getFunctionType();
   Type *FSHLRetTy = UMulFuncTy->getReturnType();
   const std::string FuncName = lowerLLVMIntrinsicName(UMulIntrinsic);
for (Instruction &I : BB) {
         Changed |= lowerIntrinsicToFunction(II);
       else if (II->getIntrinsicID() == Intrinsic::fshl ||
                II->getIntrinsicID() == Intrinsic::fshr) {
         lowerFunnelShifts(II);
         Changed = true;
       } else if (II->getIntrinsicID() == Intrinsic::umul_with_overflow) {
         lowerUMulWithOverflow(II);
         Changed = true;
+      } else if (II->getIntrinsicID() == Intrinsic::assume ||
+                 II->getIntrinsicID() == Intrinsic::expect) {
+        lowerExpectAssume(II);
+        Changed = true;
       }
     }
   }
   return Changed;
 }
 
 // Returns F if aggregate argument/return types are not present or cloned F
 // function with the types replaced by i32 types. The change in types is

llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp

 //===-- SPIRVSubtarget.cpp - SPIR-V Subtarget Information ------*- 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 implements the SPIR-V specific subclass of TargetSubtargetInfo.
 //
 //===----------------------------------------------------------------------===//
 
 #include "SPIRVSubtarget.h"
+#include "MCTargetDesc/SPIRVBaseInfo.h"
 #include "SPIRV.h"
 #include "SPIRVGlobalRegistry.h"
 #include "SPIRVLegalizerInfo.h"
 #include "SPIRVRegisterBankInfo.h"
 #include "SPIRVTargetMachine.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/TargetParser/Host.h"
 
void SPIRVSubtarget::initAvailableExtensions() {
   if (!isOpenCLEnv())
     return;
   // A default extension for testing.
   // FIXME: This should be changed when we can select extensions through a
   // command line flag.
   AvailableExtensions.insert(
       SPIRV::Extension::SPV_KHR_no_integer_wrap_decoration);
   AvailableExtensions.insert(SPIRV::Extension::SPV_INTEL_optnone);
+  AvailableExtensions.insert(SPIRV::Extension::SPV_KHR_expect_assume);
 }
 
 // TODO: use command line args for this rather than just defaults.
 // Must have called initAvailableExtensions first.
 void SPIRVSubtarget::initAvailableExtInstSets() {
   AvailableExtInstSets.clear();
   if (!isOpenCLEnv())
     AvailableExtInstSets.insert(SPIRV::InstructionSet::GLSL_std_450);

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

 defm SPV_INTEL_fpga_loop_controls : ExtensionOperand<51>;
 defm SPV_EXT_fragment_shader_interlock : ExtensionOperand<52>;
 defm SPV_NV_shader_sm_builtins : ExtensionOperand<53>;
 defm SPV_KHR_shader_clock : ExtensionOperand<54>;
 defm SPV_INTEL_unstructured_loop_controls : ExtensionOperand<55>;
 defm SPV_EXT_demote_to_helper_invocation : ExtensionOperand<56>;
 defm SPV_INTEL_fpga_reg : ExtensionOperand<57>;
 defm SPV_INTEL_optnone : ExtensionOperand<58>;
+defm SPV_KHR_expect_assume : ExtensionOperand<60>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
 // SymbolicOperand entries with string mnemonics, versioning, extensions, and
 // capabilities.
 //===----------------------------------------------------------------------===//
 
 def Capability : GenericEnum, Operand<i32> {

 defm ImageFootprintNV : CapabilityOperand<5282, 0, 0, [], []>;
 defm FragmentBarycentricNV : CapabilityOperand<5284, 0, 0, [], []>;
 defm ComputeDerivativeGroupQuadsNV : CapabilityOperand<5288, 0, 0, [], []>;
 defm ComputeDerivativeGroupLinearNV : CapabilityOperand<5350, 0, 0, [], []>;
 defm FragmentDensityEXT : CapabilityOperand<5291, 0, 0, [], [Shader]>;
 defm PhysicalStorageBufferAddressesEXT : CapabilityOperand<5347, 0, 0, [], [Shader]>;
 defm CooperativeMatrixNV : CapabilityOperand<5357, 0, 0, [], [Shader]>;
 defm OptNoneINTEL : CapabilityOperand<6094, 0, 0, [SPV_INTEL_optnone], []>;
+defm ExpectAssumeKHR : CapabilityOperand<5629, 0, 0, [SPV_KHR_expect_assume], []>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define SourceLanguage enum values and at the same time
 // SymbolicOperand entries.
 //===----------------------------------------------------------------------===//
 
 def SourceLanguage : GenericEnum, Operand<i32> {
   let FilterClass = "SourceLanguage";

llvm/test/CodeGen/SPIRV/assume.ll

+; RUN: llc -mtriple=spirv32-unknown-unknown < %s | FileCheck %s
+; RUN: llc -mtriple=spirv64-unknown-unknown < %s | FileCheck %s
+
+;      CHECK: OpCapability ExpectAssumeKHR
+; CHECK-NEXT: OpExtension "SPV_KHR_expect_assume"
+
+declare void @llvm.assume(i1)
+
+;  CHECK-DAG: %9 = OpIEqual %5 %6 %7
+; CHECK-NEXT: OpAssumeTrueKHR %9
+define void @assumeeq(i32 %x, i32 %y) {
+    %cmp = icmp eq i32 %x, %y
+    call void @llvm.assume(i1 %cmp)
+    ret void
+}
 No newline at end of file

llvm/test/CodeGen/SPIRV/expect.ll

+; RUN: llc -mtriple=spirv32-unknown-unknown < %s | FileCheck %s
+; RUN: llc -mtriple=spirv64-unknown-unknown < %s | FileCheck %s
+
+;      CHECK: OpCapability ExpectAssumeKHR
+; CHECK-NEXT: OpExtension "SPV_KHR_expect_assume"
+
+declare i32 @llvm.expect.i32(i32, i32) 
+declare i32 @getOne()
+
+; CHECK-DAG: %2 = OpTypeInt 32 0
+; CHECK-DAG: %6 = OpFunctionParameter %2
+; CHECK-DAG: %9 = OpIMul %2 %6 %8
+; CHECK-DAG: %10 = OpExpectKHR %2 %9 %6
+
+define i32 @test(i32 %x) {
+  %one = call i32 @getOne()
+  %val = mul i32 %x, %one
+  %v = call i32 @llvm.expect.i32(i32 %val, i32 %x)
+  ret i32 %v
+}