Fix VectorExtract&Insert Dynamic with sycl half translation (#1387)

This is a patch to fix forward translation of VectorExtractDynamic with sret
argument of SYCL half type & fix reverse translation of VectorInsertDynamic
with component of SYCL half type by adding extra instructions.

Author: KornevNikita

lib/SPIRV/SPIRVInternal.h

@@ -662,6 +662,8 @@ bool isOCLImageType(llvm::Type *Ty, StringRef *Name = nullptr);
 ///   type name as spirv.BaseTyName.Postfixes.
 bool isSPIRVType(llvm::Type *Ty, StringRef BaseTyName, StringRef *Postfix = 0);
 
+bool isSYCLHalfType(llvm::Type *Ty);
+
 /// Decorate a function name as __spirv_{Name}_
 std::string decorateSPIRVFunction(const std::string &S);
 
@@ -766,6 +768,19 @@ void mutateFunction(
     BuiltinFuncMangleInfo *Mangle = nullptr, AttributeList *Attrs = nullptr,
     bool TakeName = true);
 
+/// Mutate function by change the arguments & the return type.
+/// \param ArgMutate mutates the function arguments.
+/// \param RetMutate mutates the function return value.
+/// \param TakeName Take the original function's name if a new function with
+///   different type needs to be created.
+void mutateFunction(
+    Function *F,
+    std::function<std::string(CallInst *, std::vector<Value *> &, Type *&RetTy)>
+        ArgMutate,
+    std::function<Instruction *(CallInst *)> RetMutate,
+    BuiltinFuncMangleInfo *Mangle = nullptr, AttributeList *Attrs = nullptr,
+    bool TakeName = true);
+
 /// Add a call instruction at \p Pos.
 CallInst *addCallInst(Module *M, StringRef FuncName, Type *RetTy,
                       ArrayRef<Value *> Args, AttributeList *Attrs,

lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -109,6 +109,20 @@ class SPIRVRegularizeLLVMBase {
   void lowerUMulWithOverflow(IntrinsicInst *UMulIntrinsic);
   void buildUMulWithOverflowFunc(Function *UMulFunc);
 
+  // For some cases Clang emits VectorExtractDynamic as:
+  // void @_Z28__spirv_VectorExtractDynamic(<Ty>* sret(<Ty>), jointMatrix, idx);
+  // Instead of:
+  // <Ty> @_Z28__spirv_VectorExtractDynamic(JointMatrix, Idx);
+  // And VectorInsertDynamic as:
+  // @_Z27__spirv_VectorInsertDynamic(jointMatrix, <Ty>* byval(<Ty>), idx);
+  // Instead of:
+  // @_Z27__spirv_VectorInsertDynamic(jointMatrix, <Ty>, idx)
+  // Need to add additional GEP, store and load instructions and mutate called
+  // function to avoid translation failures
+  void expandSYCLHalfUsing(Module *M);
+  void expandVEDWithSYCLHalfSRetArg(Function *F);
+  void expandVIDWithSYCLHalfByValComp(Function *F);
+
   static std::string lowerLLVMIntrinsicName(IntrinsicInst *II);
   void adaptStructTypes(StructType *ST);
   static char ID;
@@ -326,6 +340,83 @@ void SPIRVRegularizeLLVMBase::lowerUMulWithOverflow(
   UMulIntrinsic->setCalledFunction(UMulFunc);
 }
 
+void SPIRVRegularizeLLVMBase::expandVEDWithSYCLHalfSRetArg(Function *F) {
+  auto Attrs = F->getAttributes();
+  Attrs = Attrs.removeParamAttribute(F->getContext(), 0, Attribute::StructRet);
+  std::string Name = F->getName().str();
+  CallInst *OldCall = nullptr;
+  mutateFunction(
+      F,
+      [=, &OldCall](CallInst *CI, std::vector<Value *> &Args, Type *&RetTy) {
+        Args.erase(Args.begin());
+        auto *SRetPtrTy = cast<PointerType>(CI->getOperand(0)->getType());
+        auto *ET = SRetPtrTy->getPointerElementType();
+        RetTy = cast<StructType>(ET)->getElementType(0);
+        OldCall = CI;
+        return Name;
+      },
+      [=, &OldCall](CallInst *NewCI) {
+        IRBuilder<> Builder(OldCall);
+        auto *SRetPtrTy = cast<PointerType>(OldCall->getOperand(0)->getType());
+        auto *ET = SRetPtrTy->getPointerElementType();
+        Value *Target = Builder.CreateStructGEP(ET, OldCall->getOperand(0), 0);
+        return Builder.CreateStore(NewCI, Target);
+      },
+      nullptr, &Attrs, true);
+}
+
+void SPIRVRegularizeLLVMBase::expandVIDWithSYCLHalfByValComp(Function *F) {
+  auto Attrs = F->getAttributes();
+  Attrs = Attrs.removeParamAttribute(F->getContext(), 1, Attribute::ByVal);
+  std::string Name = F->getName().str();
+  mutateFunction(
+      F,
+      [=](CallInst *CI, std::vector<Value *> &Args) {
+        auto *CompPtrTy = cast<PointerType>(CI->getOperand(1)->getType());
+        auto *ET = CompPtrTy->getPointerElementType();
+        Type *HalfTy = cast<StructType>(ET)->getElementType(0);
+        IRBuilder<> Builder(CI);
+        auto *Target = Builder.CreateStructGEP(ET, CI->getOperand(1), 0);
+        Args[1] = Builder.CreateLoad(HalfTy, Target);
+        return Name;
+      },
+      nullptr, &Attrs, true);
+}
+
+void SPIRVRegularizeLLVMBase::expandSYCLHalfUsing(Module *M) {
+  std::vector<Function *> ToExpandVEDWithSYCLHalfSRetArg;
+  std::vector<Function *> ToExpandVIDWithSYCLHalfByValComp;
+
+  for (auto &F : *M) {
+    if (F.getName().startswith("_Z28__spirv_VectorExtractDynamic") &&
+        F.hasStructRetAttr()) {
+      auto *SRetPtrTy = cast<PointerType>(F.getArg(0)->getType());
+      if (isSYCLHalfType(SRetPtrTy->getPointerElementType()))
+        ToExpandVEDWithSYCLHalfSRetArg.push_back(&F);
+      else
+        llvm_unreachable("The return type of the VectorExtractDynamic "
+                         "instruction cannot be a structure other than SYCL "
+                         "half.");
+    }
+    if (F.getName().startswith("_Z27__spirv_VectorInsertDynamic") &&
+        F.getArg(1)->getType()->isPointerTy()) {
+      auto *CompPtrTy = cast<PointerType>(F.getArg(1)->getType());
+      auto *ET = CompPtrTy->getPointerElementType();
+      if (isSYCLHalfType(ET))
+        ToExpandVIDWithSYCLHalfByValComp.push_back(&F);
+      else
+        llvm_unreachable("The component argument type of an "
+                         "VectorInsertDynamic instruction can't be a "
+                         "structure other than SYCL half.");
+    }
+  }
+
+  for (auto *F : ToExpandVEDWithSYCLHalfSRetArg)
+    expandVEDWithSYCLHalfSRetArg(F);
+  for (auto *F : ToExpandVIDWithSYCLHalfByValComp)
+    expandVIDWithSYCLHalfByValComp(F);
+}
+
 void SPIRVRegularizeLLVMBase::adaptStructTypes(StructType *ST) {
   if (!ST->hasName())
     return;
@@ -440,6 +531,7 @@ bool SPIRVRegularizeLLVMBase::regularize() {
   eraseUselessFunctions(M);
   lowerFuncPtr(M);
   addKernelEntryPoint(M);
+  expandSYCLHalfUsing(M);
 
   for (auto I = M->begin(), E = M->end(); I != E;) {
     Function *F = &(*I++);

lib/SPIRV/SPIRVUtil.cpp

@@ -310,6 +310,21 @@ bool isSPIRVType(llvm::Type *Ty, StringRef BaseTyName, StringRef *Postfix) {
   return false;
 }
 
+bool isSYCLHalfType(llvm::Type *Ty) {
+  if (auto *ST = dyn_cast<StructType>(Ty)) {
+    if (!ST->hasName())
+      return false;
+    StringRef Name = ST->getName();
+    Name.consume_front("class.");
+    if ((Name.startswith("cl::sycl::") ||
+         Name.startswith("__sycl_internal::")) &&
+        Name.endswith("::half")) {
+      return true;
+    }
+  }
+  return false;
+}
+
 Function *getOrCreateFunction(Module *M, Type *RetTy, ArrayRef<Type *> ArgTypes,
                               StringRef Name, BuiltinFuncMangleInfo *Mangle,
                               AttributeList *Attrs, bool TakeName) {
@@ -344,6 +359,8 @@ Function *getOrCreateFunction(Module *M, Type *RetTy, ArrayRef<Type *> ArgTypes,
     }
     LLVM_DEBUG(dbgs() << "[getOrCreateFunction] ";
                if (F) dbgs() << *F << " => "; dbgs() << *NewF << '\n';);
+    if (F)
+      NewF->setDSOLocal(F->isDSOLocal());
     F = NewF;
     F->setCallingConv(CallingConv::SPIR_FUNC);
     if (Attrs)
@@ -703,6 +720,21 @@ void mutateFunction(
     F->eraseFromParent();
 }
 
+void mutateFunction(
+    Function *F,
+    std::function<std::string(CallInst *, std::vector<Value *> &, Type *&RetTy)>
+        ArgMutate,
+    std::function<Instruction *(CallInst *)> RetMutate,
+    BuiltinFuncMangleInfo *Mangle, AttributeList *Attrs, bool TakeName) {
+  auto *M = F->getParent();
+  for (auto I = F->user_begin(), E = F->user_end(); I != E;) {
+    if (auto *CI = dyn_cast<CallInst>(*I++))
+      mutateCallInst(M, CI, ArgMutate, RetMutate, Mangle, Attrs, TakeName);
+  }
+  if (F->use_empty())
+    F->eraseFromParent();
+}
+
 CallInst *mutateCallInstSPIRV(
     Module *M, CallInst *CI,
     std::function<std::string(CallInst *, std::vector<Value *> &)> ArgMutate,