Fix type scavenger for variable arguments and multiple-uses-of-types cases. (#1606)

Original commit:
KhronosGroup/SPIRV-LLVM-Translator@401d124

Author: jcranmer-intel

llvm-spirv/lib/SPIRV/SPIRVTypeScavenger.cpp

@@ -430,10 +430,14 @@ void SPIRVTypeScavenger::correctUseTypes(Instruction &I) {
     // arguments we pass in to the argument requirements of the function.
     if (Function *F = CB->getCalledFunction()) {
       for (Use &U : CB->args()) {
+        // If we're calling a var-arg method, we have more operands than the
+        // function has parameters. Bail out if we hit that point.
+        unsigned ArgNo = CB->getArgOperandNo(&U);
+        if (ArgNo >= F->arg_size())
+          break;
         if (U->getType()->isPointerTy())
           PointerOperands.emplace_back(
-              U.getOperandNo(),
-              computePointerElementType(F->getArg(CB->getArgOperandNo(&U))));
+              U.getOperandNo(), computePointerElementType(F->getArg(ArgNo)));
       }
     }
   }
@@ -475,6 +479,16 @@ void SPIRVTypeScavenger::correctUseTypes(Instruction &I) {
       U.set(CastedValue);
     };
 
+    // This handles the scenario where a deferred type gets resolved to a fixed
+    // type during handling of this instruction, and another operand is using
+    // the same deferred type later in the instruction.
+    auto ReplaceTypeInOperands = [&](DeducedType From, DeducedType To) {
+      for (auto &ReplacePair : PointerOperands) {
+        if (ReplacePair.second == From)
+          ReplacePair.second = To;
+      }
+    };
+
     if (isa<Value *>(UsedTy)) {
       // When the use is of an indirect-pointer type, insert a bitcast to the
       // use type only for this use. This prevents indirect pointers from
@@ -491,15 +505,18 @@ void SPIRVTypeScavenger::correctUseTypes(Instruction &I) {
         // Source type is fixed, use type is deferred: set the deferred type to
         // the fixed type.
         fixType(*DeferredUseTy, FixedTy);
+        ReplaceTypeInOperands(DeferredUseTy, FixedTy);
       }
     } else if (auto *DeferredTy = dyn_cast<DeferredType *>(SourceTy)) {
       if (auto *FixedUseTy = dyn_cast<Type *>(UsedTy)) {
         // Source type is fixed, use type is deferred: set the deferred type to
         // the fixed type.
         fixType(*DeferredTy, FixedUseTy);
+        ReplaceTypeInOperands(DeferredTy, FixedUseTy);
       } else if (auto *DeferredUseTy = dyn_cast<DeferredType *>(UsedTy)) {
         // If they're both deferred, merge the two types together.
         mergeType(DeferredTy, DeferredUseTy);
+        ReplaceTypeInOperands(DeferredUseTy, DeferredTy);
       }
     }
   }

llvm-spirv/test/type-scavenger/multiple-uses.ll

@@ -0,0 +1,47 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -spirv-text -o - | FileCheck %s
+; RUN: llvm-spirv %t.bc -o %t.spv
+; RUN: spirv-val %t.spv
+
+; This test checks that, when we fix a deferred type to a known value in the
+; type scavenger, we correctly also handle replacing other types that are used
+; in the same instruction.
+
+target datalayout = "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir-unknown-unknown"
+
+; CHECK: 4 TypeInt [[INT:[0-9]+]] 32 0
+; CHECK: 4 TypePointer [[INTPTR:[0-9]+]] 7 [[INT]]
+
+; Function Attrs: nounwind
+define spir_kernel void @foo() {
+; CHECK: 4 Variable [[INTPTR]] [[IPTR:[0-9]+]] 7
+; CHECK: 4 Bitcast [[INTPTR]] [[UPTR:[0-9]+]] [[IPTR]]
+; CHECK: 4 Bitcast [[INTPTR]] [[UPTR2:[0-9]+]] [[IPTR]]
+entry:
+  %iptr = alloca i32, align 4
+  %uptr = bitcast ptr %iptr to ptr
+  %uptr2 = bitcast ptr %iptr to ptr
+  br i1 false, label %a, label %b
+
+a:
+; CHECK: 2 Label [[A:[0-9]+]]
+  br i1 false, label %c, label %d
+
+b:
+; CHECK: 2 Label [[B:[0-9]+]]
+  br label %block
+
+c:
+; CHECK: 2 Label [[C:[0-9]+]]
+  br label %block
+
+d:
+; CHECK: 2 Label [[D:[0-9]+]]
+  br label %block
+
+block:
+; CHECK: 9 Phi [[INTPTR]] {{[0-9]+}} [[UPTR]] [[B]] [[IPTR]] [[C]] [[UPTR2]] [[D]]
+  %val = phi ptr [ %uptr, %b ], [ %iptr, %c ], [ %uptr2, %d ]
+  ret void
+}

llvm-spirv/test/type-scavenger/varargs.ll

@@ -0,0 +1,24 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -spirv-text -o - | FileCheck %s
+; RUN: llvm-spirv %t.bc -o %t.spv
+
+target datalayout = "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir-unknown-unknown"
+
+@.str = internal unnamed_addr addrspace(2) constant [11 x i8] c"Value: %p\0A\00", align 1
+
+
+; CHECK: 4 TypeInt [[INT:[0-9]+]] 32 0
+; CHECK: 4 TypePointer [[INTPTR:[0-9]+]] 7 [[INT]]
+; CHECK: 5 Variable {{[0-9]+}} [[STR:[0-9]+]] 0
+; CHECK: 4 Variable [[INTPTR]] [[IPTR:[0-9]+]] 7
+; CHECK: 7 ExtInst [[INT]] {{[0-9]+}} {{[0-9]+}} printf [[STR]] [[IPTR]]
+
+; Function Attrs: nounwind
+define spir_kernel void @foo() {
+  %iptr = alloca i32, align 4
+  %res = call spir_func i32 (ptr addrspace(2), ...) @_Z18__spirv_ocl_printfPU3AS2c(ptr addrspace(2) @.str, ptr %iptr)
+  ret void
+}
+
+declare spir_func i32 @_Z18__spirv_ocl_printfPU3AS2c(ptr addrspace(2), ...)