Extend bool arg of shift operations

This is a patch to regularize lshr instruction with an i1 argument.
According to the SPIR-V specification OpShiftRightLogical operands
should be of integer type.

Co-authored-by: Alexey Sachkov <alexey.sachkov@intel.com>

Original commit:
KhronosGroup/SPIRV-LLVM-Translator@c65f190

Author: KornevNikita

llvm-spirv/lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -117,6 +117,14 @@ class SPIRVRegularizeLLVMBase {
   void expandVEDWithSYCLTypeSRetArg(Function *F);
   void expandVIDWithSYCLTypeByValComp(Function *F);
 
+  // According to the specification, the operands of a shift instruction must be
+  // a scalar/vector of integer. When LLVM-IR contains a shift instruction with
+  // i1 operands, they are treated as a bool. We need to extend them to i32 to
+  // comply with the specification. For example: "%shift = lshr i1 0, 1";
+  // The bit instruction should be changed to the extended version
+  // "%shift = lshr i32 0, 1" so the args are treated as int operands.
+  Value *extendBitInstBoolArg(Instruction *OldInst);
+
   static std::string lowerLLVMIntrinsicName(IntrinsicInst *II);
   void adaptStructTypes(StructType *ST);
   static char ID;
@@ -412,6 +420,31 @@ void SPIRVRegularizeLLVMBase::expandSYCLTypeUsing(Module *M) {
     expandVIDWithSYCLTypeByValComp(F);
 }
 
+Value *SPIRVRegularizeLLVMBase::extendBitInstBoolArg(Instruction *II) {
+  IRBuilder<> Builder(II);
+  auto *ArgTy = II->getOperand(0)->getType();
+  Type *NewArgType = nullptr;
+  if (ArgTy->isIntegerTy()) {
+    NewArgType = Builder.getInt32Ty();
+  } else if (ArgTy->isVectorTy() &&
+             cast<VectorType>(ArgTy)->getElementType()->isIntegerTy()) {
+    unsigned NumElements = cast<FixedVectorType>(ArgTy)->getNumElements();
+    NewArgType = VectorType::get(Builder.getInt32Ty(), NumElements, false);
+  } else {
+    llvm_unreachable("Unexpected type");
+  }
+  auto *NewBase = Builder.CreateZExt(II->getOperand(0), NewArgType);
+  auto *NewShift = Builder.CreateZExt(II->getOperand(1), NewArgType);
+  switch (II->getOpcode()) {
+  case Instruction::LShr:
+    return Builder.CreateLShr(NewBase, NewShift);
+  case Instruction::Shl:
+    return Builder.CreateShl(NewBase, NewShift);
+  default:
+    return II;
+  }
+}
+
 void SPIRVRegularizeLLVMBase::adaptStructTypes(StructType *ST) {
   if (!ST->hasName())
     return;
@@ -556,6 +589,21 @@ bool SPIRVRegularizeLLVMBase::regularize() {
           }
         }
 
+        // Translator treats i1 as boolean, but bit instructions take
+        // a scalar/vector integers, so we have to extend such arguments
+        if (II.isLogicalShift() &&
+            II.getOperand(0)->getType()->isIntOrIntVectorTy(1)) {
+          auto *NewInst = extendBitInstBoolArg(&II);
+          for (auto *U : II.users()) {
+            if (cast<Instruction>(U)->getOpcode() == Instruction::ZExt) {
+              U->dropAllReferences();
+              U->replaceAllUsesWith(NewInst);
+              ToErase.push_back(cast<Instruction>(U));
+            }
+          }
+          ToErase.push_back(&II);
+        }
+
         // Remove optimization info not supported by SPIRV
         if (auto BO = dyn_cast<BinaryOperator>(&II)) {
           if (isa<PossiblyExactOperator>(BO) && BO->isExact())

llvm-spirv/test/ExtendBitBoolArg.ll

@@ -0,0 +1,46 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv -s %t.bc -o %t.regulzarized.bc
+; RUN: llvm-dis %t.regulzarized.bc -o %t.regulzarized.ll
+; RUN: FileCheck < %t.regulzarized.ll %s
+
+; Translation cycle should be successfull:
+; RUN: llvm-spirv %t.regulzarized.bc -o %t.spv
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+
+; CHECK: %[[#Base:]] = load i1, i1 addrspace(4)*{{.*}}, align 1
+; CHECK: %[[#LoadShift:]] = load i32, i32 addrspace(4)*{{.*}} align 4
+; CHECK: %[[#AndShift:]] = and i32 %[[#LoadShift]], 1
+; CHECK: %[[#CmpShift:]] = icmp ne i32 %[[#AndShift]], 0
+; CHECK: %[[#ExtBase:]] = select i1 %[[#Base]], i32 1, i32 0
+; CHECK: %[[#ExtShift:]] = select i1 %[[#CmpShift]], i32 1, i32 0
+; CHECK: %[[#LSHR:]] = lshr i32 %[[#ExtBase]], %[[#ExtShift]]
+; CHECK: and i32 %[[#LSHR]], 1
+
+; CHECK: %[[#ExtVecBase:]] = select <2 x i1> %vec1, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+; CHECK: %[[#ExtVecShift:]] = select <2 x i1> %vec2, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+; CHECK: lshr <2 x i32> %[[#ExtVecBase]], %[[#ExtVecShift]]
+
+; ModuleID = 'source.bc'
+source_filename = "source.cpp"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+%"class.ac" = type { i1 }
+
+; Function Attrs: convergent mustprogress norecurse nounwind
+define linkonce_odr dso_local spir_func void @foo(<2 x i1> %vec1, <2 x i1> %vec2) align 2 {
+  %1 = alloca %"class.ac" addrspace(4)*, align 8
+  %2 = alloca i32, align 4
+  %3 = addrspacecast %"class.ac" addrspace(4)** %1 to %"class.ac" addrspace(4)* addrspace(4)*
+  %4 = addrspacecast i32* %2 to i32 addrspace(4)*
+  %5 = load %"class.ac" addrspace(4)*, %"class.ac" addrspace(4)* addrspace(4)* %3, align 8
+  %6 = getelementptr inbounds %"class.ac", %"class.ac" addrspace(4)* %5, i32 0, i32 0
+  %7 = load i1, i1 addrspace(4)* %6, align 1
+  %8 = load i32, i32 addrspace(4)* %4, align 4
+  %9 = trunc i32 %8 to i1
+  %10 = lshr i1 %7, %9
+  %11 = zext i1 %10 to i32
+  %12 = and i32 %11, 1
+  %13 = lshr <2 x i1> %vec1, %vec2
+  ret void
+}