[SPIRV] fix several issues in builds with expensive checks

[iliya-diyachkov]

Summary

The patch fixes "Virtual register does not match instruction constraint" and partly "Illegal virtual register for instruction" fails in SPIRV backend builds with LLVM_ENABLE_EXPENSIVE_CHECKS enabled. As a result the number of passed LIT tests with the enabled checks is increased from 59 to 120.

Also, support for ndrange_1D/ndrange_2D/ndrange_3D builtins is placed in a separate function.

[arsenm]

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

392

I'd expect the class to never be set at this point (if not, I'd expect it to be always set)

[iliya-diyachkov]

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

392

18 LIT tests will fail if I remove this setRegClass(). On other hand, I found calls with register classes already set here in 27 LIT tests.
So I suppose it makes sense to leave this check and setting.

[zuban32]

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

392

@arsenm could you pls elaborate? Is it supposed to be set later (if so, where)? I thought CallLowering being a part of IRTranslator is a perfect place to set reg classes as vregs have just been created

[arsenm]

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

392

If the instruction had a register constraint it should have been set to begin with. You are fixing up invalid mir after the fact

[mpaszkowski]

@iliya-diyachkov I checked the half / vstore_half_rtn test manually and it seems to work just fine. Looks like we just got one flaky run for some reason.

[iliya-diyachkov]

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

392

I removed this fixing and added the required setRegClasss in SPIRVBuiltins.cpp. Now we set reg classes only when replacing calls with machine instructions.

[iliya-diyachkov]

In D144897#4168668, @mpaszkowski wrote:

@iliya-diyachkov I checked the half / vstore_half_rtn test manually and it seems to work just fine. Looks like we just got one flaky run for some reason.

Thanks, Michal. However we need to retest the patch since changes in buildScopeReg and buildMemSemanticsReg can affect some tests.

[mpaszkowski]

Let's wait for the CTS testing results before pushing, but the patch LGTM! Thank you @iliya-diyachkov!

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

const FixedVectorType *LLVMVectorType =
         cast<FixedVectorType>(GR->getTypeForSPIRVType(BoolType));
     Type = LLT::vector(LLVMVectorType->getElementCount(), 1);
   } else {
     Type = LLT::scalar(1);
   }
 
   Register ResultRegister =
       MIRBuilder.getMRI()->createGenericVirtualRegister(Type);
+  MIRBuilder.getMRI()->setRegClass(ResultRegister, &SPIRV::IDRegClass);
   GR->assignSPIRVTypeToVReg(BoolType, ResultRegister, MIRBuilder.getMF());
   return std::make_tuple(ResultRegister, BoolType);
 }
 
 /// Helper function for building either a vector or scalar select instruction
 /// depending on the expected \p ResultType.
 static bool buildSelectInst(MachineIRBuilder &MIRBuilder,
                             Register ReturnRegister, Register SourceRegister,

 static Register buildConstantIntReg(uint64_t Val, MachineIRBuilder &MIRBuilder,
                                     SPIRVGlobalRegistry *GR,
                                     unsigned BitWidth = 32) {
   SPIRVType *IntType = GR->getOrCreateSPIRVIntegerType(BitWidth, MIRBuilder);
   return GR->buildConstantInt(Val, MIRBuilder, IntType);
 }
 
 static Register buildScopeReg(Register CLScopeRegister,
+                              SPIRV::Scope::Scope Scope,
                               MachineIRBuilder &MIRBuilder,
                               SPIRVGlobalRegistry *GR,
-                              const MachineRegisterInfo *MRI) {
+                              MachineRegisterInfo *MRI) {
+  if (CLScopeRegister.isValid()) {
     auto CLScope =
         static_cast<SPIRV::CLMemoryScope>(getIConstVal(CLScopeRegister, MRI));
-  SPIRV::Scope::Scope Scope = getSPIRVScope(CLScope);
+    Scope = getSPIRVScope(CLScope);
 
-  if (CLScope == static_cast<unsigned>(Scope))
+    if (CLScope == static_cast<unsigned>(Scope)) {
+      MRI->setRegClass(CLScopeRegister, &SPIRV::IDRegClass);
       return CLScopeRegister;
-
+    }
+  }
   return buildConstantIntReg(Scope, MIRBuilder, GR);
 }
 
 static Register buildMemSemanticsReg(Register SemanticsRegister,
-                                     Register PtrRegister,
-                                     const MachineRegisterInfo *MRI,
+                                     Register PtrRegister, unsigned &Semantics,
+                                     MachineIRBuilder &MIRBuilder,
                                      SPIRVGlobalRegistry *GR) {
+  if (SemanticsRegister.isValid()) {
+    MachineRegisterInfo *MRI = MIRBuilder.getMRI();
     std::memory_order Order =
         static_cast<std::memory_order>(getIConstVal(SemanticsRegister, MRI));
-  unsigned Semantics =
+    Semantics =
         getSPIRVMemSemantics(Order) |
         getMemSemanticsForStorageClass(GR->getPointerStorageClass(PtrRegister));
 
-  if (Order == Semantics)
+    if (Order == Semantics) {
+      MRI->setRegClass(SemanticsRegister, &SPIRV::IDRegClass);
       return SemanticsRegister;
-
-  return Register();
+    }
+  }
+  return buildConstantIntReg(Semantics, MIRBuilder, GR);
 }
 
 /// Helper function for translating atomic init to OpStore.
 static bool buildAtomicInitInst(const SPIRV::IncomingCall *Call,
                                 MachineIRBuilder &MIRBuilder) {
   assert(Call->Arguments.size() == 2 &&
          "Need 2 arguments for atomic init translation");
-
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
   MIRBuilder.buildInstr(SPIRV::OpStore)
       .addUse(Call->Arguments[0])
       .addUse(Call->Arguments[1]);
   return true;
 }
 
 /// Helper function for building an atomic load instruction.
 static bool buildAtomicLoadInst(const SPIRV::IncomingCall *Call,
                                 MachineIRBuilder &MIRBuilder,
                                 SPIRVGlobalRegistry *GR) {
   Register PtrRegister = Call->Arguments[0];
+  MIRBuilder.getMRI()->setRegClass(PtrRegister, &SPIRV::IDRegClass);
   // TODO: if true insert call to __translate_ocl_memory_sccope before
   // OpAtomicLoad and the function implementation. We can use Translator's
   // output for transcoding/atomic_explicit_arguments.cl as an example.
   Register ScopeRegister;
-  if (Call->Arguments.size() > 1)
+  if (Call->Arguments.size() > 1) {
     ScopeRegister = Call->Arguments[1];
-  else
+    MIRBuilder.getMRI()->setRegClass(ScopeRegister, &SPIRV::IDRegClass);
+  } else
     ScopeRegister = buildConstantIntReg(SPIRV::Scope::Device, MIRBuilder, GR);
 
   Register MemSemanticsReg;
   if (Call->Arguments.size() > 2) {
     // TODO: Insert call to __translate_ocl_memory_order before OpAtomicLoad.
     MemSemanticsReg = Call->Arguments[2];
+    MIRBuilder.getMRI()->setRegClass(MemSemanticsReg, &SPIRV::IDRegClass);
   } else {
     int Semantics =
         SPIRV::MemorySemantics::SequentiallyConsistent |
         getMemSemanticsForStorageClass(GR->getPointerStorageClass(PtrRegister));
     MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
   }
 
   MIRBuilder.buildInstr(SPIRV::OpAtomicLoad)
       .addDef(Call->ReturnRegister)
       .addUse(GR->getSPIRVTypeID(Call->ReturnType))
       .addUse(PtrRegister)
       .addUse(ScopeRegister)
       .addUse(MemSemanticsReg);
   return true;
 }
 
 /// Helper function for building an atomic store instruction.
 static bool buildAtomicStoreInst(const SPIRV::IncomingCall *Call,
                                  MachineIRBuilder &MIRBuilder,
                                  SPIRVGlobalRegistry *GR) {
   Register ScopeRegister =
       buildConstantIntReg(SPIRV::Scope::Device, MIRBuilder, GR);
   Register PtrRegister = Call->Arguments[0];
+  MIRBuilder.getMRI()->setRegClass(PtrRegister, &SPIRV::IDRegClass);
   int Semantics =
       SPIRV::MemorySemantics::SequentiallyConsistent |
       getMemSemanticsForStorageClass(GR->getPointerStorageClass(PtrRegister));
   Register MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
-
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
   MIRBuilder.buildInstr(SPIRV::OpAtomicStore)
       .addUse(PtrRegister)
       .addUse(ScopeRegister)
       .addUse(MemSemanticsReg)
       .addUse(Call->Arguments[1]);
   return true;
 }
 
 /// Helper function for building an atomic compare-exchange instruction.
 static bool buildAtomicCompareExchangeInst(const SPIRV::IncomingCall *Call,
                                            MachineIRBuilder &MIRBuilder,
                                            SPIRVGlobalRegistry *GR) {
   const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
   unsigned Opcode =
       SPIRV::lookupNativeBuiltin(Builtin->Name, Builtin->Set)->Opcode;
   bool IsCmpxchg = Call->Builtin->Name.contains("cmpxchg");
   MachineRegisterInfo *MRI = MIRBuilder.getMRI();
 
   Register ObjectPtr = Call->Arguments[0];   // Pointer (volatile A *object.)
   Register ExpectedArg = Call->Arguments[1]; // Comparator (C* expected).
   Register Desired = Call->Arguments[2];     // Value (C Desired).
+  MRI->setRegClass(ObjectPtr, &SPIRV::IDRegClass);
+  MRI->setRegClass(ExpectedArg, &SPIRV::IDRegClass);
+  MRI->setRegClass(Desired, &SPIRV::IDRegClass);
   SPIRVType *SpvDesiredTy = GR->getSPIRVTypeForVReg(Desired);
   LLT DesiredLLT = MRI->getType(Desired);
 
   assert(GR->getSPIRVTypeForVReg(ObjectPtr)->getOpcode() ==
          SPIRV::OpTypePointer);
   unsigned ExpectedType = GR->getSPIRVTypeForVReg(ExpectedArg)->getOpcode();
   assert(IsCmpxchg ? ExpectedType == SPIRV::OpTypeInt
                    : ExpectedType == SPIRV::OpTypePointer);
if (Call->Arguments.size() >= 4) {
     auto MemOrdNeq =
         static_cast<std::memory_order>(getIConstVal(Call->Arguments[4], MRI));
     MemSemEqual = getSPIRVMemSemantics(MemOrdEq) | MemSemStorage;
     MemSemUnequal = getSPIRVMemSemantics(MemOrdNeq) | MemSemStorage;
     if (MemOrdEq == MemSemEqual)
       MemSemEqualReg = Call->Arguments[3];
     if (MemOrdNeq == MemSemEqual)
       MemSemUnequalReg = Call->Arguments[4];
+    MRI->setRegClass(Call->Arguments[3], &SPIRV::IDRegClass);
+    MRI->setRegClass(Call->Arguments[4], &SPIRV::IDRegClass);
   }
   if (!MemSemEqualReg.isValid())
     MemSemEqualReg = buildConstantIntReg(MemSemEqual, MIRBuilder, GR);
   if (!MemSemUnequalReg.isValid())
     MemSemUnequalReg = buildConstantIntReg(MemSemUnequal, MIRBuilder, GR);
 
   Register ScopeReg;
   auto Scope = IsCmpxchg ? SPIRV::Scope::Workgroup : SPIRV::Scope::Device;
   if (Call->Arguments.size() >= 6) {
     assert(Call->Arguments.size() == 6 &&
            "Extra args for explicit atomic cmpxchg");
     auto ClScope = static_cast<SPIRV::CLMemoryScope>(
         getIConstVal(Call->Arguments[5], MRI));
     Scope = getSPIRVScope(ClScope);
     if (ClScope == static_cast<unsigned>(Scope))
       ScopeReg = Call->Arguments[5];
+    MRI->setRegClass(Call->Arguments[5], &SPIRV::IDRegClass);
   }
   if (!ScopeReg.isValid())
     ScopeReg = buildConstantIntReg(Scope, MIRBuilder, GR);
 
   Register Expected = IsCmpxchg
                           ? ExpectedArg
                           : buildLoadInst(SpvDesiredTy, ExpectedArg, MIRBuilder,
                                           GR, LLT::scalar(32));
   MRI->setType(Expected, DesiredLLT);
   Register Tmp = !IsCmpxchg ? MRI->createGenericVirtualRegister(DesiredLLT)
                             : Call->ReturnRegister;
+  if (!MRI->getRegClassOrNull(Tmp))
+    MRI->setRegClass(Tmp, &SPIRV::IDRegClass);
   GR->assignSPIRVTypeToVReg(SpvDesiredTy, Tmp, MIRBuilder.getMF());
 
   SPIRVType *IntTy = GR->getOrCreateSPIRVIntegerType(32, MIRBuilder);
   MIRBuilder.buildInstr(Opcode)
       .addDef(Tmp)
       .addUse(GR->getSPIRVTypeID(IntTy))
       .addUse(ObjectPtr)
       .addUse(ScopeReg)
       .addUse(MemSemEqualReg)
       .addUse(MemSemUnequalReg)
       .addUse(Desired)
       .addUse(Expected);
   if (!IsCmpxchg) {
     MIRBuilder.buildInstr(SPIRV::OpStore).addUse(ExpectedArg).addUse(Tmp);
     MIRBuilder.buildICmp(CmpInst::ICMP_EQ, Call->ReturnRegister, Tmp, Expected);
   }
   return true;
 }
 
 /// Helper function for building an atomic load instruction.
 static bool buildAtomicRMWInst(const SPIRV::IncomingCall *Call, unsigned Opcode,
                                MachineIRBuilder &MIRBuilder,
                                SPIRVGlobalRegistry *GR) {
-  const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
-  SPIRV::Scope::Scope Scope = SPIRV::Scope::Workgroup;
-  Register ScopeRegister;
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  Register ScopeRegister =
+      Call->Arguments.size() >= 4 ? Call->Arguments[3] : Register();
 
-  if (Call->Arguments.size() >= 4) {
-    assert(Call->Arguments.size() == 4 &&
-           "Too many args for explicit atomic RMW");
-    ScopeRegister = buildScopeReg(Call->Arguments[3], MIRBuilder, GR, MRI);
-  }
-
-  if (!ScopeRegister.isValid())
-    ScopeRegister = buildConstantIntReg(Scope, MIRBuilder, GR);
+  assert(Call->Arguments.size() <= 4 &&
+  ScopeRegister = buildScopeReg(ScopeRegister, SPIRV::Scope::Workgroup,
+                                MIRBuilder, GR, MRI);
 
   Register PtrRegister = Call->Arguments[0];
   unsigned Semantics = SPIRV::MemorySemantics::None;
-  Register MemSemanticsReg;
-
-  if (Call->Arguments.size() >= 3)
-    MemSemanticsReg =
-        buildMemSemanticsReg(Call->Arguments[2], PtrRegister, MRI, GR);
-
-  if (!MemSemanticsReg.isValid())
-    MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
-
+  MRI->setRegClass(PtrRegister, &SPIRV::IDRegClass);
+  Register MemSemanticsReg =
+      Call->Arguments.size() >= 3 ? Call->Arguments[2] : Register();
+  MemSemanticsReg = buildMemSemanticsReg(MemSemanticsReg, PtrRegister,
+                                         Semantics, MIRBuilder, GR);
+  MRI->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
   MIRBuilder.buildInstr(Opcode)
       .addDef(Call->ReturnRegister)
       .addUse(GR->getSPIRVTypeID(Call->ReturnType))
       .addUse(PtrRegister)
       .addUse(ScopeRegister)
       .addUse(MemSemanticsReg)
       .addUse(Call->Arguments[1]);
   return true;
 }
 
 /// Helper function for building atomic flag instructions (e.g.
 /// OpAtomicFlagTestAndSet).
 static bool buildAtomicFlagInst(const SPIRV::IncomingCall *Call,
                                 unsigned Opcode, MachineIRBuilder &MIRBuilder,
                                 SPIRVGlobalRegistry *GR) {
-  const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
-
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   Register PtrRegister = Call->Arguments[0];
   unsigned Semantics = SPIRV::MemorySemantics::SequentiallyConsistent;
-  Register MemSemanticsReg;
-
-  if (Call->Arguments.size() >= 2)
-    MemSemanticsReg =
-        buildMemSemanticsReg(Call->Arguments[1], PtrRegister, MRI, GR);
-
-  if (!MemSemanticsReg.isValid())
-    MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
+  Register MemSemanticsReg =
+      Call->Arguments.size() >= 2 ? Call->Arguments[1] : Register();
+  MemSemanticsReg = buildMemSemanticsReg(MemSemanticsReg, PtrRegister,
+                                         Semantics, MIRBuilder, GR);
 
   assert((Opcode != SPIRV::OpAtomicFlagClear ||
           (Semantics != SPIRV::MemorySemantics::Acquire &&
            Semantics != SPIRV::MemorySemantics::AcquireRelease)) &&
          "Invalid memory order argument!");
 
-  SPIRV::Scope::Scope Scope = SPIRV::Scope::Device;
-  Register ScopeRegister;
-
-  if (Call->Arguments.size() >= 3)
-    ScopeRegister = buildScopeReg(Call->Arguments[2], MIRBuilder, GR, MRI);
-
-  if (!ScopeRegister.isValid())
-    ScopeRegister = buildConstantIntReg(Scope, MIRBuilder, GR);
+  Register ScopeRegister =
+      Call->Arguments.size() >= 3 ? Call->Arguments[2] : Register();
+  ScopeRegister =
+      buildScopeReg(ScopeRegister, SPIRV::Scope::Device, MIRBuilder, GR, MRI);
 
   auto MIB = MIRBuilder.buildInstr(Opcode);
   if (Opcode == SPIRV::OpAtomicFlagTestAndSet)
     MIB.addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Call->ReturnType));
 
   MIB.addUse(PtrRegister).addUse(ScopeRegister).addUse(MemSemanticsReg);
   return true;
 }
 
 /// Helper function for building barriers, i.e., memory/control ordering
 /// operations.
 static bool buildBarrierInst(const SPIRV::IncomingCall *Call, unsigned Opcode,
                              MachineIRBuilder &MIRBuilder,
                              SPIRVGlobalRegistry *GR) {
-  const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   unsigned MemFlags = getIConstVal(Call->Arguments[0], MRI);
   unsigned MemSemantics = SPIRV::MemorySemantics::None;
 
   if (MemFlags & SPIRV::CLK_LOCAL_MEM_FENCE)
     MemSemantics |= SPIRV::MemorySemantics::WorkgroupMemory;
 
   if (MemFlags & SPIRV::CLK_GLOBAL_MEM_FENCE)
     MemSemantics |= SPIRV::MemorySemantics::CrossWorkgroupMemory;
 
   if (MemFlags & SPIRV::CLK_IMAGE_MEM_FENCE)
     MemSemantics |= SPIRV::MemorySemantics::ImageMemory;
 
   if (Opcode == SPIRV::OpMemoryBarrier) {
     std::memory_order MemOrder =
         static_cast<std::memory_order>(getIConstVal(Call->Arguments[1], MRI));
     MemSemantics = getSPIRVMemSemantics(MemOrder) | MemSemantics;
   } else {
     MemSemantics |= SPIRV::MemorySemantics::SequentiallyConsistent;
   }
 
   Register MemSemanticsReg;
-  if (MemFlags == MemSemantics)
+  if (MemFlags == MemSemantics) {
     MemSemanticsReg = Call->Arguments[0];
-  else
+    MRI->setRegClass(MemSemanticsReg, &SPIRV::IDRegClass);
+  } else
     MemSemanticsReg = buildConstantIntReg(MemSemantics, MIRBuilder, GR);
 
   Register ScopeReg;
   SPIRV::Scope::Scope Scope = SPIRV::Scope::Workgroup;
   SPIRV::Scope::Scope MemScope = Scope;
   if (Call->Arguments.size() >= 2) {
     assert(
         ((Opcode != SPIRV::OpMemoryBarrier && Call->Arguments.size() == 2) ||
          (Opcode == SPIRV::OpMemoryBarrier && Call->Arguments.size() == 3)) &&
         "Extra args for explicitly scoped barrier");
     Register ScopeArg = (Opcode == SPIRV::OpMemoryBarrier) ? Call->Arguments[2]
                                                            : Call->Arguments[1];
     SPIRV::CLMemoryScope CLScope =
         static_cast<SPIRV::CLMemoryScope>(getIConstVal(ScopeArg, MRI));
     MemScope = getSPIRVScope(CLScope);
     if (!(MemFlags & SPIRV::CLK_LOCAL_MEM_FENCE) ||
         (Opcode == SPIRV::OpMemoryBarrier))
       Scope = MemScope;
 
-    if (CLScope == static_cast<unsigned>(Scope))
+    if (CLScope == static_cast<unsigned>(Scope)) {
       ScopeReg = Call->Arguments[1];
+      MRI->setRegClass(ScopeReg, &SPIRV::IDRegClass);
+    }
   }
 
   if (!ScopeReg.isValid())
     ScopeReg = buildConstantIntReg(Scope, MIRBuilder, GR);
 
   auto MIB = MIRBuilder.buildInstr(Opcode).addUse(ScopeReg);
   if (Opcode != SPIRV::OpMemoryBarrier)
     MIB.addUse(buildConstantIntReg(MemScope, MIRBuilder, GR));

 }
 
 static bool generateGroupInst(const SPIRV::IncomingCall *Call,
                               MachineIRBuilder &MIRBuilder,
                               SPIRVGlobalRegistry *GR) {
   const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
   const SPIRV::GroupBuiltin *GroupBuiltin =
       SPIRV::lookupGroupBuiltin(Builtin->Name);
-  const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   Register Arg0;
   if (GroupBuiltin->HasBoolArg) {
     Register ConstRegister = Call->Arguments[0];
     auto ArgInstruction = getDefInstrMaybeConstant(ConstRegister, MRI);
     // TODO: support non-constant bool values.
     assert(ArgInstruction->getOpcode() == TargetOpcode::G_CONSTANT &&
            "Only constant bool value args are supported");
     if (GR->getSPIRVTypeForVReg(Call->Arguments[0])->getOpcode() !=
auto MIB = MIRBuilder.buildInstr(GroupBuiltin->Opcode)
                  .addDef(GroupResultRegister)
                  .addUse(GR->getSPIRVTypeID(GroupResultType))
                  .addUse(ScopeRegister);
 
   if (!GroupBuiltin->NoGroupOperation)
     MIB.addImm(GroupBuiltin->GroupOperation);
   if (Call->Arguments.size() > 0) {
     MIB.addUse(Arg0.isValid() ? Arg0 : Call->Arguments[0]);
-    for (unsigned i = 1; i < Call->Arguments.size(); i++)
+    MRI->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+    for (unsigned i = 1; i < Call->Arguments.size(); i++) {
       MIB.addUse(Call->Arguments[i]);
+      MRI->setRegClass(Call->Arguments[i], &SPIRV::IDRegClass);
+    }
   }
 
   // Build select instruction.
   if (HasBoolReturnTy)
     buildSelectInst(MIRBuilder, Call->ReturnRegister, GroupResultRegister,
                     Call->ReturnType, GR);
   return true;
 }
static bool genWorkgroupQuery(const SPIRV::IncomingCall *Call,
 
   // If the index is constant, we can statically determine if it is in range.
   bool IsConstantIndex =
       IndexInstruction->getOpcode() == TargetOpcode::G_CONSTANT;
 
   // If it's out of range (max dimension is 3), we can just return the constant
   // default value (0 or 1 depending on which query function).
   if (IsConstantIndex && getIConstVal(IndexRegister, MRI) >= 3) {
-    Register defaultReg = Call->ReturnRegister;
+    Register DefaultReg = Call->ReturnRegister;
     if (PointerSize != ResultWidth) {
-      defaultReg = MRI->createGenericVirtualRegister(LLT::scalar(PointerSize));
-      GR->assignSPIRVTypeToVReg(PointerSizeType, defaultReg,
+      DefaultReg = MRI->createGenericVirtualRegister(LLT::scalar(PointerSize));
+      MRI->setRegClass(DefaultReg, &SPIRV::IDRegClass);
+      GR->assignSPIRVTypeToVReg(PointerSizeType, DefaultReg,
                                 MIRBuilder.getMF());
-      ToTruncate = defaultReg;
+      ToTruncate = DefaultReg;
     }
     auto NewRegister =
         GR->buildConstantInt(DefaultValue, MIRBuilder, PointerSizeType);
-    MIRBuilder.buildCopy(defaultReg, NewRegister);
+    MIRBuilder.buildCopy(DefaultReg, NewRegister);
   } else { // If it could be in range, we need to load from the given builtin.
     auto Vec3Ty =
         GR->getOrCreateSPIRVVectorType(PointerSizeType, 3, MIRBuilder);
     Register LoadedVector =
         buildBuiltinVariableLoad(MIRBuilder, Vec3Ty, GR, BuiltinValue,
                                  LLT::fixed_vector(3, PointerSize));
     // Set up the vreg to extract the result to (possibly a new temporary one).
     Register Extracted = Call->ReturnRegister;
     if (!IsConstantIndex || PointerSize != ResultWidth) {
       Extracted = MRI->createGenericVirtualRegister(LLT::scalar(PointerSize));
+      MRI->setRegClass(Extracted, &SPIRV::IDRegClass);
       GR->assignSPIRVTypeToVReg(PointerSizeType, Extracted, MIRBuilder.getMF());
     }
     // Use Intrinsic::spv_extractelt so dynamic vs static extraction is
     // handled later: extr = spv_extractelt LoadedVector, IndexRegister.
     MachineInstrBuilder ExtractInst = MIRBuilder.buildIntrinsic(
         Intrinsic::spv_extractelt, ArrayRef<Register>{Extracted}, true);
     ExtractInst.addUse(LoadedVector).addUse(IndexRegister);
 
     // If the index is dynamic, need check if it's < 3, and then use a select.
     if (!IsConstantIndex) {
       insertAssignInstr(Extracted, nullptr, PointerSizeType, GR, MIRBuilder,
                         *MRI);
 
       auto IndexType = GR->getSPIRVTypeForVReg(IndexRegister);
       auto BoolType = GR->getOrCreateSPIRVBoolType(MIRBuilder);
 
       Register CompareRegister =
           MRI->createGenericVirtualRegister(LLT::scalar(1));
+      MRI->setRegClass(CompareRegister, &SPIRV::IDRegClass);
       GR->assignSPIRVTypeToVReg(BoolType, CompareRegister, MIRBuilder.getMF());
 
       // Use G_ICMP to check if idxVReg < 3.
       MIRBuilder.buildICmp(CmpInst::ICMP_ULT, CompareRegister, IndexRegister,
                            GR->buildConstantInt(3, MIRBuilder, IndexType));
 
       // Get constant for the default value (0 or 1 depending on which
       // function).
       Register DefaultRegister =
           GR->buildConstantInt(DefaultValue, MIRBuilder, PointerSizeType);
 
       // Get a register for the selection result (possibly a new temporary one).
       Register SelectionResult = Call->ReturnRegister;
       if (PointerSize != ResultWidth) {
         SelectionResult =
             MRI->createGenericVirtualRegister(LLT::scalar(PointerSize));
+        MRI->setRegClass(SelectionResult, &SPIRV::IDRegClass);
         GR->assignSPIRVTypeToVReg(PointerSizeType, SelectionResult,
                                   MIRBuilder.getMF());
       }
       // Create the final G_SELECT to return the extracted value or the default.
       MIRBuilder.buildSelect(SelectionResult, CompareRegister, Extracted,
                              DefaultRegister);
       ToTruncate = SelectionResult;
     } else {
static bool generateImageSizeQueryInst(const SPIRV::IncomingCall *Call,
   // Get the actual number of query result/size components.
   SPIRVType *ImgType = GR->getSPIRVTypeForVReg(Call->Arguments[0]);
   unsigned NumActualRetComponents = getNumSizeComponents(ImgType);
   Register QueryResult = Call->ReturnRegister;
   SPIRVType *QueryResultType = Call->ReturnType;
   if (NumExpectedRetComponents != NumActualRetComponents) {
     QueryResult = MIRBuilder.getMRI()->createGenericVirtualRegister(
         LLT::fixed_vector(NumActualRetComponents, 32));
+    MIRBuilder.getMRI()->setRegClass(QueryResult, &SPIRV::IDRegClass);
     SPIRVType *IntTy = GR->getOrCreateSPIRVIntegerType(32, MIRBuilder);
     QueryResultType = GR->getOrCreateSPIRVVectorType(
         IntTy, NumActualRetComponents, MIRBuilder);
     GR->assignSPIRVTypeToVReg(QueryResultType, QueryResult, MIRBuilder.getMF());
   }
   bool IsDimBuf = ImgType->getOperand(2).getImm() == SPIRV::Dim::DIM_Buffer;
   unsigned Opcode =
       IsDimBuf ? SPIRV::OpImageQuerySize : SPIRV::OpImageQuerySizeLod;
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
   auto MIB = MIRBuilder.buildInstr(Opcode)
                  .addDef(QueryResult)
                  .addUse(GR->getSPIRVTypeID(QueryResultType))
                  .addUse(Call->Arguments[0]);
   if (!IsDimBuf)
     MIB.addUse(buildConstantIntReg(0, MIRBuilder, GR)); // Lod id.
   if (NumExpectedRetComponents == NumActualRetComponents)
     return true;
assert(Call->ReturnType->getOpcode() == SPIRV::OpTypeInt &&
          "Image samples query result must be of int type!");
 
   // Lookup the instruction opcode in the TableGen records.
   const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
   unsigned Opcode =
       SPIRV::lookupNativeBuiltin(Builtin->Name, Builtin->Set)->Opcode;
 
   Register Image = Call->Arguments[0];
+  MIRBuilder.getMRI()->setRegClass(Image, &SPIRV::IDRegClass);
   SPIRV::Dim::Dim ImageDimensionality = static_cast<SPIRV::Dim::Dim>(
       GR->getSPIRVTypeForVReg(Image)->getOperand(2).getImm());
 
   switch (Opcode) {
   case SPIRV::OpImageQuerySamples:
     assert(ImageDimensionality == SPIRV::Dim::DIM_2D &&
            "Image must be of 2D dimensionality");
     break;

 }
 
 static bool generateReadImageInst(const StringRef DemangledCall,
                                   const SPIRV::IncomingCall *Call,
                                   MachineIRBuilder &MIRBuilder,
                                   SPIRVGlobalRegistry *GR) {
   Register Image = Call->Arguments[0];
   MachineRegisterInfo *MRI = MIRBuilder.getMRI();
-
-  if (DemangledCall.contains_insensitive("ocl_sampler")) {
+  MRI->setRegClass(Image, &SPIRV::IDRegClass);
+  MRI->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
+  bool HasOclSampler = DemangledCall.contains_insensitive("ocl_sampler");
+  bool HasMsaa = DemangledCall.contains_insensitive("msaa");
+  if (HasOclSampler || HasMsaa)
+    MRI->setRegClass(Call->Arguments[2], &SPIRV::IDRegClass);
+  if (HasOclSampler) {
     Register Sampler = Call->Arguments[1];
 
     if (!GR->isScalarOfType(Sampler, SPIRV::OpTypeSampler) &&
         getDefInstrMaybeConstant(Sampler, MRI)->getOperand(1).isCImm()) {
       uint64_t SamplerMask = getIConstVal(Sampler, MRI);
       Sampler = GR->buildConstantSampler(
           Register(), getSamplerAddressingModeFromBitmask(SamplerMask),
           getSamplerParamFromBitmask(SamplerMask),
if (HasOclSampler) {
     bool NeedsExtraction = false;
     if (TempType->getOpcode() != SPIRV::OpTypeVector) {
       TempType =
           GR->getOrCreateSPIRVVectorType(Call->ReturnType, 4, MIRBuilder);
       NeedsExtraction = true;
     }
     LLT LLType = LLT::scalar(GR->getScalarOrVectorBitWidth(TempType));
     Register TempRegister = MRI->createGenericVirtualRegister(LLType);
+    MRI->setRegClass(TempRegister, &SPIRV::IDRegClass);
     GR->assignSPIRVTypeToVReg(TempType, TempRegister, MIRBuilder.getMF());
 
     MIRBuilder.buildInstr(SPIRV::OpImageSampleExplicitLod)
         .addDef(NeedsExtraction ? TempRegister : Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(TempType))
         .addUse(SampledImage)
         .addUse(Call->Arguments[2]) // Coordinate.
         .addImm(SPIRV::ImageOperand::Lod)
         .addUse(Lod);
 
     if (NeedsExtraction)
       MIRBuilder.buildInstr(SPIRV::OpCompositeExtract)
           .addDef(Call->ReturnRegister)
           .addUse(GR->getSPIRVTypeID(Call->ReturnType))
           .addUse(TempRegister)
           .addImm(0);
-  } else if (DemangledCall.contains_insensitive("msaa")) {
+  } else if (HasMsaa) {
     MIRBuilder.buildInstr(SPIRV::OpImageRead)
         .addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Call->ReturnType))
         .addUse(Image)
         .addUse(Call->Arguments[1]) // Coordinate.
         .addImm(SPIRV::ImageOperand::Sample)
         .addUse(Call->Arguments[2]);
   } else {
     MIRBuilder.buildInstr(SPIRV::OpImageRead)
         .addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Call->ReturnType))
         .addUse(Image)
         .addUse(Call->Arguments[1]); // Coordinate.
   }
   return true;
 }
 
 static bool generateWriteImageInst(const SPIRV::IncomingCall *Call,
                                    MachineIRBuilder &MIRBuilder,
                                    SPIRVGlobalRegistry *GR) {
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
+  MIRBuilder.getMRI()->setRegClass(Call->Arguments[2], &SPIRV::IDRegClass);
   MIRBuilder.buildInstr(SPIRV::OpImageWrite)
       .addUse(Call->Arguments[0])  // Image.
       .addUse(Call->Arguments[1])  // Coordinate.
       .addUse(Call->Arguments[2]); // Texel.
   return true;
 }
 
 static bool generateSampleImageInst(const StringRef DemangledCall,
                                     const SPIRV::IncomingCall *Call,
                                     MachineIRBuilder &MIRBuilder,
                                     SPIRVGlobalRegistry *GR) {
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   if (Call->Builtin->Name.contains_insensitive(
           "__translate_sampler_initializer")) {
     // Build sampler literal.
-    uint64_t Bitmask = getIConstVal(Call->Arguments[0], MIRBuilder.getMRI());
+    uint64_t Bitmask = getIConstVal(Call->Arguments[0], MRI);
     Register Sampler = GR->buildConstantSampler(
         Call->ReturnRegister, getSamplerAddressingModeFromBitmask(Bitmask),
         getSamplerParamFromBitmask(Bitmask),
         getSamplerFilterModeFromBitmask(Bitmask), MIRBuilder, Call->ReturnType);
     return Sampler.isValid();
   } else if (Call->Builtin->Name.contains_insensitive("__spirv_SampledImage")) {
     // Create OpSampledImage.
     Register Image = Call->Arguments[0];
     SPIRVType *ImageType = GR->getSPIRVTypeForVReg(Image);
     SPIRVType *SampledImageType =
         GR->getOrCreateOpTypeSampledImage(ImageType, MIRBuilder);
     Register SampledImage =
         Call->ReturnRegister.isValid()
             ? Call->ReturnRegister
-            : MIRBuilder.getMRI()->createVirtualRegister(&SPIRV::IDRegClass);
+            : MRI->createVirtualRegister(&SPIRV::IDRegClass);
     MIRBuilder.buildInstr(SPIRV::OpSampledImage)
         .addDef(SampledImage)
         .addUse(GR->getSPIRVTypeID(SampledImageType))
         .addUse(Image)
         .addUse(Call->Arguments[1]); // Sampler.
     return true;
   } else if (Call->Builtin->Name.contains_insensitive(
                  "__spirv_ImageSampleExplicitLod")) {
     // Sample an image using an explicit level of detail.
     std::string ReturnType = DemangledCall.str();
     if (DemangledCall.contains("_R")) {
       ReturnType = ReturnType.substr(ReturnType.find("_R") + 2);
       ReturnType = ReturnType.substr(0, ReturnType.find('('));
     }
     SPIRVType *Type = GR->getOrCreateSPIRVTypeByName(ReturnType, MIRBuilder);
+    MRI->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+    MRI->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
+    MRI->setRegClass(Call->Arguments[3], &SPIRV::IDRegClass);
+
     MIRBuilder.buildInstr(SPIRV::OpImageSampleExplicitLod)
         .addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Type))
         .addUse(Call->Arguments[0]) // Image.
         .addUse(Call->Arguments[1]) // Coordinate.
         .addImm(SPIRV::ImageOperand::Lod)
         .addUse(Call->Arguments[3]);
     return true;
for (unsigned i = 0; i < Call->Arguments.size(); i++)
       MIB.addUse(Call->Arguments[i]);
     return true;
   }
   default:
     return false;
   }
 }
 
+static bool buildNDRange(const SPIRV::IncomingCall *Call,
+                         MachineIRBuilder &MIRBuilder,
+                         SPIRVGlobalRegistry *GR) {
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  MRI->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+  SPIRVType *PtrType = GR->getSPIRVTypeForVReg(Call->Arguments[0]);
+  assert(PtrType->getOpcode() == SPIRV::OpTypePointer &&
+         PtrType->getOperand(2).isReg());
+  Register TypeReg = PtrType->getOperand(2).getReg();
+  SPIRVType *StructType = GR->getSPIRVTypeForVReg(TypeReg);
+  MachineFunction &MF = MIRBuilder.getMF();
+  Register TmpReg = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+  GR->assignSPIRVTypeToVReg(StructType, TmpReg, MF);
+  // Skip the first arg, it's the destination pointer. OpBuildNDRange takes
+  // three other arguments, so pass zero constant on absence.
+  unsigned NumArgs = Call->Arguments.size();
+  assert(NumArgs >= 2);
+  Register GlobalWorkSize = Call->Arguments[NumArgs < 4 ? 1 : 2];
+  MRI->setRegClass(GlobalWorkSize, &SPIRV::IDRegClass);
+  Register LocalWorkSize =
+      NumArgs == 2 ? Register(0) : Call->Arguments[NumArgs < 4 ? 2 : 3];
+  if (LocalWorkSize.isValid())
+    MRI->setRegClass(LocalWorkSize, &SPIRV::IDRegClass);
+  Register GlobalWorkOffset = NumArgs <= 3 ? Register(0) : Call->Arguments[1];
+  if (GlobalWorkOffset.isValid())
+    MRI->setRegClass(GlobalWorkOffset, &SPIRV::IDRegClass);
+  if (NumArgs < 4) {
+    Register Const;
+    SPIRVType *SpvTy = GR->getSPIRVTypeForVReg(GlobalWorkSize);
+    if (SpvTy->getOpcode() == SPIRV::OpTypePointer) {
+      MachineInstr *DefInstr = MRI->getUniqueVRegDef(GlobalWorkSize);
+      assert(DefInstr && isSpvIntrinsic(*DefInstr, Intrinsic::spv_gep) &&
+             DefInstr->getOperand(3).isReg());
+      Register GWSPtr = DefInstr->getOperand(3).getReg();
+      if (!MRI->getRegClassOrNull(GWSPtr))
+        MRI->setRegClass(GWSPtr, &SPIRV::IDRegClass);
+      // TODO: Maybe simplify generation of the type of the fields.
+      unsigned Size = Call->Builtin->Name.equals("ndrange_3D") ? 3 : 2;
+      unsigned BitWidth = GR->getPointerSize() == 64 ? 64 : 32;
+      Type *BaseTy = IntegerType::get(MF.getFunction().getContext(), BitWidth);
+      Type *FieldTy = ArrayType::get(BaseTy, Size);
+      SPIRVType *SpvFieldTy = GR->getOrCreateSPIRVType(FieldTy, MIRBuilder);
+      GlobalWorkSize = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+      GR->assignSPIRVTypeToVReg(SpvFieldTy, GlobalWorkSize, MF);
+      MIRBuilder.buildInstr(SPIRV::OpLoad)
+          .addDef(GlobalWorkSize)
+          .addUse(GR->getSPIRVTypeID(SpvFieldTy))
+          .addUse(GWSPtr);
+      Const = GR->getOrCreateConsIntArray(0, MIRBuilder, SpvFieldTy);
+    } else {
+      Const = GR->buildConstantInt(0, MIRBuilder, SpvTy);
+    }
+    if (!LocalWorkSize.isValid())
+      LocalWorkSize = Const;
+    if (!GlobalWorkOffset.isValid())
+      GlobalWorkOffset = Const;
+  }
+  assert(LocalWorkSize.isValid() && GlobalWorkOffset.isValid());
+  MIRBuilder.buildInstr(SPIRV::OpBuildNDRange)
+      .addDef(TmpReg)
+      .addUse(TypeReg)
+      .addUse(GlobalWorkSize)
+      .addUse(LocalWorkSize)
+      .addUse(GlobalWorkOffset);
+  return MIRBuilder.buildInstr(SPIRV::OpStore)
+      .addUse(Call->Arguments[0])
+      .addUse(TmpReg);
+}
+
 static MachineInstr *getBlockStructInstr(Register ParamReg,
                                          MachineRegisterInfo *MRI) {
   // We expect the following sequence of instructions:
   //   %0:_(pN) = G_INTRINSIC_W_SIDE_EFFECTS intrinsic(@llvm.spv.alloca)
   //   or       = G_GLOBAL_VALUE @block_literal_global
   //   %1:_(pN) = G_INTRINSIC_W_SIDE_EFFECTS intrinsic(@llvm.spv.bitcast), %0
   //   %2:_(p4) = G_ADDRSPACE_CAST %1:_(pN)
   MachineInstr *MI = MRI->getUniqueVRegDef(ParamReg);
if (Call->Builtin->Name.find("_varargs") != StringRef::npos) {
     assert(LocalSizeTy && "Local size type is expected");
     const uint64_t LocalSizeNum =
         cast<ArrayType>(LocalSizeTy)->getNumElements();
     unsigned SC = storageClassToAddressSpace(SPIRV::StorageClass::Generic);
     const LLT LLType = LLT::pointer(SC, GR->getPointerSize());
     const SPIRVType *PointerSizeTy = GR->getOrCreateSPIRVPointerType(
         Int32Ty, MIRBuilder, SPIRV::StorageClass::Function);
     for (unsigned I = 0; I < LocalSizeNum; ++I) {
-      Register Reg =
-          MIRBuilder.getMRI()->createVirtualRegister(&SPIRV::IDRegClass);
-      MIRBuilder.getMRI()->setType(Reg, LLType);
+      Register Reg = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+      MRI->setType(Reg, LLType);
       GR->assignSPIRVTypeToVReg(PointerSizeTy, Reg, MIRBuilder.getMF());
       auto GEPInst = MIRBuilder.buildIntrinsic(Intrinsic::spv_gep,
                                                ArrayRef<Register>{Reg}, true);
       GEPInst
           .addImm(GepMI->getOperand(2).getImm())          // In bound.
           .addUse(ArrayMI->getOperand(0).getReg())        // Alloca.
           .addUse(buildConstantIntReg(0, MIRBuilder, GR)) // Indices.
           .addUse(buildConstantIntReg(I, MIRBuilder, GR));
static bool generateEnqueueInst(const SPIRV::IncomingCall *Call,
   // Lookup the instruction opcode in the TableGen records.
   const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
   unsigned Opcode =
       SPIRV::lookupNativeBuiltin(Builtin->Name, Builtin->Set)->Opcode;
 
   switch (Opcode) {
   case SPIRV::OpRetainEvent:
   case SPIRV::OpReleaseEvent:
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
     return MIRBuilder.buildInstr(Opcode).addUse(Call->Arguments[0]);
   case SPIRV::OpCreateUserEvent:
   case SPIRV::OpGetDefaultQueue:
     return MIRBuilder.buildInstr(Opcode)
         .addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Call->ReturnType));
   case SPIRV::OpIsValidEvent:
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
     return MIRBuilder.buildInstr(Opcode)
         .addDef(Call->ReturnRegister)
         .addUse(GR->getSPIRVTypeID(Call->ReturnType))
         .addUse(Call->Arguments[0]);
   case SPIRV::OpSetUserEventStatus:
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
     return MIRBuilder.buildInstr(Opcode)
         .addUse(Call->Arguments[0])
         .addUse(Call->Arguments[1]);
   case SPIRV::OpCaptureEventProfilingInfo:
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
+    MIRBuilder.getMRI()->setRegClass(Call->Arguments[2], &SPIRV::IDRegClass);
     return MIRBuilder.buildInstr(Opcode)
         .addUse(Call->Arguments[0])
         .addUse(Call->Arguments[1])
         .addUse(Call->Arguments[2]);
-  case SPIRV::OpBuildNDRange: {
-    MachineRegisterInfo *MRI = MIRBuilder.getMRI();
-    SPIRVType *PtrType = GR->getSPIRVTypeForVReg(Call->Arguments[0]);
-    assert(PtrType->getOpcode() == SPIRV::OpTypePointer &&
-           PtrType->getOperand(2).isReg());
-    Register TypeReg = PtrType->getOperand(2).getReg();
-    SPIRVType *StructType = GR->getSPIRVTypeForVReg(TypeReg);
-    Register TmpReg = MRI->createVirtualRegister(&SPIRV::IDRegClass);
-    GR->assignSPIRVTypeToVReg(StructType, TmpReg, MIRBuilder.getMF());
-    // Skip the first arg, it's the destination pointer. OpBuildNDRange takes
-    // three other arguments, so pass zero constant on absence.
-    unsigned NumArgs = Call->Arguments.size();
-    assert(NumArgs >= 2);
-    Register GlobalWorkSize = Call->Arguments[NumArgs < 4 ? 1 : 2];
-    Register LocalWorkSize =
-        NumArgs == 2 ? Register(0) : Call->Arguments[NumArgs < 4 ? 2 : 3];
-    Register GlobalWorkOffset = NumArgs <= 3 ? Register(0) : Call->Arguments[1];
-    if (NumArgs < 4) {
-      Register Const;
-      SPIRVType *SpvTy = GR->getSPIRVTypeForVReg(GlobalWorkSize);
-      if (SpvTy->getOpcode() == SPIRV::OpTypePointer) {
-        MachineInstr *DefInstr = MRI->getUniqueVRegDef(GlobalWorkSize);
-        assert(DefInstr && isSpvIntrinsic(*DefInstr, Intrinsic::spv_gep) &&
-               DefInstr->getOperand(3).isReg());
-        Register GWSPtr = DefInstr->getOperand(3).getReg();
-        // TODO: Maybe simplify generation of the type of the fields.
-        unsigned Size = Call->Builtin->Name.equals("ndrange_3D") ? 3 : 2;
-        unsigned BitWidth = GR->getPointerSize() == 64 ? 64 : 32;
-        Type *BaseTy = IntegerType::get(
-            MIRBuilder.getMF().getFunction().getContext(), BitWidth);
-        Type *FieldTy = ArrayType::get(BaseTy, Size);
-        SPIRVType *SpvFieldTy = GR->getOrCreateSPIRVType(FieldTy, MIRBuilder);
-        GlobalWorkSize = MRI->createVirtualRegister(&SPIRV::IDRegClass);
-        GR->assignSPIRVTypeToVReg(SpvFieldTy, GlobalWorkSize,
-                                  MIRBuilder.getMF());
-        MIRBuilder.buildInstr(SPIRV::OpLoad)
-            .addDef(GlobalWorkSize)
-            .addUse(GR->getSPIRVTypeID(SpvFieldTy))
-            .addUse(GWSPtr);
-        Const = GR->getOrCreateConsIntArray(0, MIRBuilder, SpvFieldTy);
-      } else {
-        Const = GR->buildConstantInt(0, MIRBuilder, SpvTy);
-      }
-      if (!LocalWorkSize.isValid())
-        LocalWorkSize = Const;
-      if (!GlobalWorkOffset.isValid())
-        GlobalWorkOffset = Const;
-    }
-    MIRBuilder.buildInstr(Opcode)
-        .addDef(TmpReg)
-        .addUse(TypeReg)
-        .addUse(GlobalWorkSize)
-        .addUse(LocalWorkSize)
-        .addUse(GlobalWorkOffset);
-    return MIRBuilder.buildInstr(SPIRV::OpStore)
-        .addUse(Call->Arguments[0])
-        .addUse(TmpReg);
-  }
+  case SPIRV::OpBuildNDRange:
+    return buildNDRange(Call, MIRBuilder, GR);
   case SPIRV::OpEnqueueKernel:
     return buildEnqueueKernel(Call, MIRBuilder, GR);
   default:
     return false;
   }
 }
 
 static bool generateAsyncCopy(const SPIRV::IncomingCall *Call,
static bool generateLoadStoreInst(const SPIRV::IncomingCall *Call,
   // Build the instruction.
   auto MIB = MIRBuilder.buildInstr(Opcode);
   if (IsLoad) {
     MIB.addDef(Call->ReturnRegister);
     MIB.addUse(GR->getSPIRVTypeID(Call->ReturnType));
   }
   // Add a pointer to the value to load/store.
   MIB.addUse(Call->Arguments[0]);
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  MRI->setRegClass(Call->Arguments[0], &SPIRV::IDRegClass);
   // Add a value to store.
-  if (!IsLoad)
+  if (!IsLoad) {
     MIB.addUse(Call->Arguments[1]);
+    MRI->setRegClass(Call->Arguments[1], &SPIRV::IDRegClass);
+  }
   // Add optional memory attributes and an alignment.
-  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   unsigned NumArgs = Call->Arguments.size();
-  if ((IsLoad && NumArgs >= 2) || NumArgs >= 3)
+  if ((IsLoad && NumArgs >= 2) || NumArgs >= 3) {
     MIB.addImm(getConstFromIntrinsic(Call->Arguments[IsLoad ? 1 : 2], MRI));
-  if ((IsLoad && NumArgs >= 3) || NumArgs >= 4)
+    MRI->setRegClass(Call->Arguments[IsLoad ? 1 : 2], &SPIRV::IDRegClass);
+  }
+  if ((IsLoad && NumArgs >= 3) || NumArgs >= 4) {
     MIB.addImm(getConstFromIntrinsic(Call->Arguments[IsLoad ? 2 : 3], MRI));
+    MRI->setRegClass(Call->Arguments[IsLoad ? 2 : 3], &SPIRV::IDRegClass);
+  }
   return true;
 }
 
 /// Lowers a builtin funtion call using the provided \p DemangledCall skeleton
 /// and external instruction \p Set.
 namespace SPIRV {
 std::optional<bool> lowerBuiltin(const StringRef DemangledCall,
                                  SPIRV::InstructionSet::InstructionSet Set,
                                  MachineIRBuilder &MIRBuilder,
                                  const Register OrigRet, const Type *OrigRetTy,
                                  const SmallVectorImpl<Register> &Args,
                                  SPIRVGlobalRegistry *GR) {
   LLVM_DEBUG(dbgs() << "Lowering builtin call: " << DemangledCall << "\n");
 
   // SPIR-V type and return register.
   Register ReturnRegister = OrigRet;
   SPIRVType *ReturnType = nullptr;
   if (OrigRetTy && !OrigRetTy->isVoidTy()) {
     ReturnType = GR->assignTypeToVReg(OrigRetTy, OrigRet, MIRBuilder);
+    if (!MIRBuilder.getMRI()->getRegClassOrNull(ReturnRegister))
+      MIRBuilder.getMRI()->setRegClass(ReturnRegister, &SPIRV::IDRegClass);
   } else if (OrigRetTy && OrigRetTy->isVoidTy()) {
     ReturnRegister = MIRBuilder.getMRI()->createVirtualRegister(&IDRegClass);
     MIRBuilder.getMRI()->setType(ReturnRegister, LLT::scalar(32));
     ReturnType = GR->assignTypeToVReg(OrigRetTy, ReturnRegister, MIRBuilder);
   }
 
   // Lookup the builtin in the TableGen records.
   std::unique_ptr<const IncomingCall> Call =

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

bool SPIRVCallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   if (Info.Callee.isGlobal()) {
     CF = dyn_cast_or_null<const Function>(Info.Callee.getGlobal());
     // TODO: support constexpr casts and indirect calls.
     if (CF == nullptr)
       return false;
     FTy = getOriginalFunctionType(*CF);
   }
 
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   Register ResVReg =
       Info.OrigRet.Regs.empty() ? Register(0) : Info.OrigRet.Regs[0];
   std::string FuncName = Info.Callee.getGlobal()->getName().str();
   std::string DemangledName = getOclOrSpirvBuiltinDemangledName(FuncName);
   const auto *ST = static_cast<const SPIRVSubtarget *>(&MF.getSubtarget());
   // TODO: check that it's OCL builtin, then apply OpenCL_std.
   if (!DemangledName.empty() && CF && CF->isDeclaration() &&
       ST->canUseExtInstSet(SPIRV::InstructionSet::OpenCL_std)) {
     const Type *OrigRetTy = Info.OrigRet.Ty;
     if (FTy)
       OrigRetTy = FTy->getReturnType();
     SmallVector<Register, 8> ArgVRegs;
     for (auto Arg : Info.OrigArgs) {
       assert(Arg.Regs.size() == 1 && "Call arg has multiple VRegs");
       ArgVRegs.push_back(Arg.Regs[0]);
       SPIRVType *SPIRVTy = GR->getOrCreateSPIRVType(Arg.Ty, MIRBuilder);
       GR->assignSPIRVTypeToVReg(SPIRVTy, Arg.Regs[0], MIRBuilder.getMF());
     }
     if (auto Res = SPIRV::lowerBuiltin(
             DemangledName, SPIRV::InstructionSet::OpenCL_std, MIRBuilder,
             ResVReg, OrigRetTy, ArgVRegs, GR))
       return *Res;
   }
   if (CF && CF->isDeclaration() &&
       !GR->find(CF, &MIRBuilder.getMF()).isValid()) {
     // Emit the type info and forward function declaration to the first MBB
     // to ensure VReg definition dependencies are valid across all MBBs.
     MachineIRBuilder FirstBlockBuilder;
     FirstBlockBuilder.setMF(MF);
     FirstBlockBuilder.setMBB(*MF.getBlockNumbered(0));
 
     SmallVector<ArrayRef<Register>, 8> VRegArgs;
     SmallVector<SmallVector<Register, 1>, 8> ToInsert;
     for (const Argument &Arg : CF->args()) {
       if (MIRBuilder.getDataLayout().getTypeStoreSize(Arg.getType()).isZero())
         continue; // Don't handle zero sized types.
-      ToInsert.push_back(
-          {MIRBuilder.getMRI()->createGenericVirtualRegister(LLT::scalar(32))});
+      Register Reg = MRI->createGenericVirtualRegister(LLT::scalar(32));
+      MRI->setRegClass(Reg, &SPIRV::IDRegClass);
+      ToInsert.push_back({Reg});
       VRegArgs.push_back(ToInsert.back());
     }
     // TODO: Reuse FunctionLoweringInfo
     FunctionLoweringInfo FuncInfo;
     lowerFormalArguments(FirstBlockBuilder, *CF, VRegArgs, FuncInfo);
   }
 
   // Make sure there's a valid return reg, even for functions returning void.

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

SPIRVGlobalRegistry::getOrCreateConstIntReg(uint64_t Val, SPIRVType *SpvType,
   bool NewInstr = false;
   // Find a constant in DT or build a new one.
   ConstantInt *CI = ConstantInt::get(const_cast<IntegerType *>(LLVMIntTy), Val);
   Register Res = DT.find(CI, CurMF);
   if (!Res.isValid()) {
     unsigned BitWidth = SpvType ? getScalarOrVectorBitWidth(SpvType) : 32;
     LLT LLTy = LLT::scalar(32);
     Res = CurMF->getRegInfo().createGenericVirtualRegister(LLTy);
+    CurMF->getRegInfo().setRegClass(Res, &SPIRV::IDRegClass);
     if (MIRBuilder)
       assignTypeToVReg(LLVMIntTy, Res, *MIRBuilder);
     else
       assignIntTypeToVReg(BitWidth, Res, *I, *TII);
     DT.add(CI, CurMF, Res);
     NewInstr = true;
   }
   return std::make_tuple(Res, CI, NewInstr);
Register SPIRVGlobalRegistry::buildConstantInt(uint64_t Val,
   // Find a constant in DT or build a new one.
   const auto ConstInt =
       ConstantInt::get(const_cast<IntegerType *>(LLVMIntTy), Val);
   Register Res = DT.find(ConstInt, &MF);
   if (!Res.isValid()) {
     unsigned BitWidth = SpvType ? getScalarOrVectorBitWidth(SpvType) : 32;
     LLT LLTy = LLT::scalar(EmitIR ? BitWidth : 32);
     Res = MF.getRegInfo().createGenericVirtualRegister(LLTy);
+    MF.getRegInfo().setRegClass(Res, &SPIRV::IDRegClass);
     assignTypeToVReg(LLVMIntTy, Res, MIRBuilder,
                      SPIRV::AccessQualifier::ReadWrite, EmitIR);
     DT.add(ConstInt, &MIRBuilder.getMF(), Res);
     if (EmitIR) {
       MIRBuilder.buildConstant(Res, *ConstInt);
     } else {
       MachineInstrBuilder MIB;
       if (Val) {
if (SpvType) {
     LLVMFPTy = IntegerType::getFloatTy(MF.getFunction().getContext());
   }
   // Find a constant in DT or build a new one.
   const auto ConstFP = ConstantFP::get(LLVMFPTy->getContext(), Val);
   Register Res = DT.find(ConstFP, &MF);
   if (!Res.isValid()) {
     unsigned BitWidth = SpvType ? getScalarOrVectorBitWidth(SpvType) : 32;
     Res = MF.getRegInfo().createGenericVirtualRegister(LLT::scalar(BitWidth));
+    MF.getRegInfo().setRegClass(Res, &SPIRV::IDRegClass);
     assignTypeToVReg(LLVMFPTy, Res, MIRBuilder);
     DT.add(ConstFP, &MF, Res);
     MIRBuilder.buildFConstant(Res, *ConstFP);
   }
   return Res;
 }
 
 Register SPIRVGlobalRegistry::getOrCreateIntCompositeOrNull(
if (!Res.isValid()) {
     // TODO: can moved below once sorting of types/consts/defs is implemented.
     Register SpvScalConst;
     if (Val)
       SpvScalConst = getOrCreateConstInt(Val, I, SpvBaseType, TII);
     // TODO: maybe use bitwidth of base type.
     LLT LLTy = LLT::scalar(32);
     Register SpvVecConst =
         CurMF->getRegInfo().createGenericVirtualRegister(LLTy);
+    CurMF->getRegInfo().setRegClass(SpvVecConst, &SPIRV::IDRegClass);
     assignSPIRVTypeToVReg(SpvType, SpvVecConst, *CurMF);
     DT.add(CA, CurMF, SpvVecConst);
     MachineInstrBuilder MIB;
     MachineBasicBlock &BB = *I.getParent();
     if (Val) {
       MIB = BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpConstantComposite))
                 .addDef(SpvVecConst)
                 .addUse(getSPIRVTypeID(SpvType));
if (!Res.isValid()) {
     if (Val || EmitIR) {
       SPIRVType *SpvBaseType =
           getOrCreateSPIRVIntegerType(BitWidth, MIRBuilder);
       SpvScalConst = buildConstantInt(Val, MIRBuilder, SpvBaseType, EmitIR);
     }
     LLT LLTy = EmitIR ? LLT::fixed_vector(ElemCnt, BitWidth) : LLT::scalar(32);
     Register SpvVecConst =
         CurMF->getRegInfo().createGenericVirtualRegister(LLTy);
+    CurMF->getRegInfo().setRegClass(SpvVecConst, &SPIRV::IDRegClass);
     assignSPIRVTypeToVReg(SpvType, SpvVecConst, *CurMF);
     DT.add(CA, CurMF, SpvVecConst);
     if (EmitIR) {
       MIRBuilder.buildSplatVector(SpvVecConst, SpvScalConst);
     } else {
       if (Val) {
         auto MIB = MIRBuilder.buildInstr(SPIRV::OpConstantComposite)
                        .addDef(SpvVecConst)
SPIRVGlobalRegistry::getOrCreateConstNullPtr(MachineIRBuilder &MIRBuilder,
   const Type *LLVMTy = getTypeForSPIRVType(SpvType);
   const PointerType *LLVMPtrTy = cast<PointerType>(LLVMTy);
   // Find a constant in DT or build a new one.
   Constant *CP = ConstantPointerNull::get(const_cast<PointerType *>(LLVMPtrTy));
   Register Res = DT.find(CP, CurMF);
   if (!Res.isValid()) {
     LLT LLTy = LLT::pointer(LLVMPtrTy->getAddressSpace(), PointerSize);
     Res = CurMF->getRegInfo().createGenericVirtualRegister(LLTy);
+    CurMF->getRegInfo().setRegClass(Res, &SPIRV::IDRegClass);
     assignSPIRVTypeToVReg(SpvType, Res, *CurMF);
     MIRBuilder.buildInstr(SPIRV::OpConstantNull)
         .addDef(Res)
         .addUse(getSPIRVTypeID(SpvType));
     DT.add(CP, CurMF, Res);
   }
   return Res;
 }
Register SPIRVGlobalRegistry::getOrCreateUndef(MachineInstr &I,
   assert(LLVMTy);
   // Find a constant in DT or build a new one.
   UndefValue *UV = UndefValue::get(const_cast<Type *>(LLVMTy));
   Register Res = DT.find(UV, CurMF);
   if (Res.isValid())
     return Res;
   LLT LLTy = LLT::scalar(32);
   Res = CurMF->getRegInfo().createGenericVirtualRegister(LLTy);
+  CurMF->getRegInfo().setRegClass(Res, &SPIRV::IDRegClass);
   assignSPIRVTypeToVReg(SpvType, Res, *CurMF);
   DT.add(UV, CurMF, Res);
 
   MachineInstrBuilder MIB;
   MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(SPIRV::OpUndef))
             .addDef(Res)
             .addUse(getSPIRVTypeID(SpvType));
   const auto &ST = CurMF->getSubtarget();
   constrainSelectedInstRegOperands(*MIB, *ST.getInstrInfo(),
                                    *ST.getRegisterInfo(), *ST.getRegBankInfo());
   return Res;
 }

llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp

for (MachineInstr &MI : MBB) {
         }
       }
     }
   }
   for (MachineInstr *MI : ToErase) {
     Register Reg = MI->getOperand(2).getReg();
     if (RegsAlreadyAddedToDT.find(MI) != RegsAlreadyAddedToDT.end())
       Reg = RegsAlreadyAddedToDT[MI];
+    auto *RC = MRI.getRegClassOrNull(MI->getOperand(0).getReg());
+    if (!MRI.getRegClassOrNull(Reg) && RC)
+      MRI.setRegClass(Reg, RC);
     MRI.replaceRegWith(MI->getOperand(0).getReg(), Reg);
     MI->eraseFromParent();
   }
   for (MachineInstr *MI : ToEraseComposites)
     MI->eraseFromParent();
 }
 
 static void foldConstantsIntoIntrinsics(MachineFunction &MF) {
Register insertAssignInstr(Register Reg, Type *Ty, SPIRVType *SpirvTy,
                            SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
                            MachineRegisterInfo &MRI) {
   MachineInstr *Def = MRI.getVRegDef(Reg);
   assert((Ty || SpirvTy) && "Either LLVM or SPIRV type is expected.");
   MIB.setInsertPt(*Def->getParent(),
                   (Def->getNextNode() ? Def->getNextNode()->getIterator()
                                       : Def->getParent()->end()));
   Register NewReg = MRI.createGenericVirtualRegister(MRI.getType(Reg));
-  if (auto *RC = MRI.getRegClassOrNull(Reg))
+  if (auto *RC = MRI.getRegClassOrNull(Reg)) {
     MRI.setRegClass(NewReg, RC);
+  } else {
+    MRI.setRegClass(NewReg, &SPIRV::IDRegClass);
+    MRI.setRegClass(Reg, &SPIRV::IDRegClass);
+  }
   SpirvTy = SpirvTy ? SpirvTy : GR->getOrCreateSPIRVType(Ty, MIB);
   GR->assignSPIRVTypeToVReg(SpirvTy, Reg, MIB.getMF());
   // This is to make it convenient for Legalizer to get the SPIRVType
   // when processing the actual MI (i.e. not pseudo one).
   GR->assignSPIRVTypeToVReg(SpirvTy, NewReg, MIB.getMF());
   // Copy MIFlags from Def to ASSIGN_TYPE instruction. It's required to keep
   // the flags after instruction selection.
   const uint16_t Flags = Def->getFlags();
   MIB.buildInstr(SPIRV::ASSIGN_TYPE)
       .addDef(Reg)
       .addUse(NewReg)
       .addUse(GR->getSPIRVTypeID(SpirvTy))
       .setMIFlags(Flags);
   Def->getOperand(0).setReg(NewReg);
-  MRI.setRegClass(Reg, &SPIRV::ANYIDRegClass);
   return NewReg;
 }
 } // namespace llvm
 
 static void generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                                  MachineIRBuilder MIB) {
   MachineRegisterInfo &MRI = MF.getRegInfo();
   SmallVector<MachineInstr *, 10> ToErase;