[SPIRV] Implement support for SPV_KHR_expect_assume (#66217)

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.

Previously https://reviews.llvm.org/D157696

Author: pmatos

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -33,4 +33,8 @@ let TargetPrefix = "spv" in {
   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>]>;
 }

llvm/lib/Target/SPIRV/SPIRVInstrInfo.td

@@ -91,6 +91,10 @@ def OpNop: SimpleOp<"OpNop", 0>;
 def OpUndef: Op<1, (outs ID:$res), (ins TYPE:$type), "$res = OpUndef $type">;
 def OpSizeOf: Op<321, (outs ID:$res), (ins TYPE:$ty, ID:$ptr), "$res = OpSizeOf $ty $ptr">;
 
+//  - 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">;
+
 // 3.42.2 Debug Instructions
 
 def OpSourceContinued: Op<2, (outs), (ins StringImm:$str, variable_ops),

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "MCTargetDesc/SPIRVMCTargetDesc.h"
 #include "SPIRV.h"
 #include "SPIRVGlobalRegistry.h"
 #include "SPIRVInstrInfo.h"
@@ -1395,6 +1396,17 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
     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");
   }

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -903,6 +903,13 @@ void addInstrRequirements(const MachineInstr &MI,
   case SPIRV::OpGroupNonUniformBallotFindMSB:
     Reqs.addCapability(SPIRV::Capability::GroupNonUniformBallot);
     break;
+  case SPIRV::OpAssumeTrueKHR:
+  case SPIRV::OpExpectKHR:
+    if (ST.canUseExtension(SPIRV::Extension::SPV_KHR_expect_assume)) {
+      Reqs.addExtension(SPIRV::Extension::SPV_KHR_expect_assume);
+      Reqs.addCapability(SPIRV::Capability::ExpectAssumeKHR);
+    }
+    break;
   default:
     break;
   }

llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp

@@ -24,6 +24,8 @@
 #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"
 
@@ -233,6 +235,32 @@ static void buildUMulWithOverflowFunc(Function *UMulFunc) {
   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
@@ -270,6 +298,10 @@ bool SPIRVPrepareFunctions::substituteIntrinsicCalls(Function *F) {
       } 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;
       }
     }
   }

llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp

@@ -11,7 +11,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "SPIRVSubtarget.h"
-#include "MCTargetDesc/SPIRVBaseInfo.h"
 #include "SPIRV.h"
 #include "SPIRVGlobalRegistry.h"
 #include "SPIRVLegalizerInfo.h"

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -451,6 +451,7 @@ defm CooperativeMatrixNV : CapabilityOperand<5357, 0, 0, [], [Shader]>;
 defm ArbitraryPrecisionIntegersINTEL : CapabilityOperand<5844, 0, 0, [SPV_INTEL_arbitrary_precision_integers], [Int8, Int16]>;
 defm OptNoneINTEL : CapabilityOperand<6094, 0, 0, [SPV_INTEL_optnone], []>;
 defm BitInstructions : CapabilityOperand<6025, 0, 0, [SPV_KHR_bit_instructions], []>;
+defm ExpectAssumeKHR : CapabilityOperand<5629, 0, 0, [SPV_KHR_expect_assume], []>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define SourceLanguage enum values and at the same time

llvm/test/CodeGen/SPIRV/assume.ll

@@ -0,0 +1,15 @@
+; 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
+}

llvm/test/CodeGen/SPIRV/expect.ll

@@ -0,0 +1,20 @@
+; 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
+}