Handle @llvm.memset.* with non-constant arguments (intel#696)

* Handle @llvm.memset.* with non-constant arguments

There is no SPIR-V counterpart for @llvm.memset.* intrinsic. Cases with constant
value and length arguments are emulated via "storing" a constant array to the
destination.

For other cases we wrap the intrinsic in @spirv.llvm_memset_* function and
expand the intrinsic to a loop via expandMemSetAsLoop() from
llvm/Transforms/Utils/LowerMemIntrinsics.h.

During reverse translation from SPIR-V to LLVM IR we can detect
@spirv.llvm_memset_* and replace it with @llvm.memset.

Signed-off-by: Alexey Sotkin <[email protected]>
Co-authored-by: Sven van Haastregt <[email protected]>

Author: 2 people

llvm-spirv/lib/SPIRV/SPIRVReader.cpp

@@ -2786,21 +2786,33 @@ Function *SPIRVToLLVM::transFunction(SPIRVFunction *BF) {
   auto IsKernel = isKernel(BF);
   auto Linkage = IsKernel ? GlobalValue::ExternalLinkage : transLinkageType(BF);
   FunctionType *FT = dyn_cast<FunctionType>(transType(BF->getFunctionType()));
-  Function *F = cast<Function>(
-      mapValue(BF, Function::Create(FT, Linkage, BF->getName(), M)));
+  std::string FuncName = BF->getName();
+  StringRef FuncNameRef(FuncName);
+  // Transform "@spirv.llvm_memset_p0i8_i32.volatile" to @llvm.memset.p0i8.i32
+  // assuming llvm.memset is supported by the device compiler. If this
+  // assumption is not safe, we should have a command line option to control
+  // this behavior.
+  if (FuncNameRef.consume_front("spirv.")) {
+    FuncNameRef.consume_back(".volatile");
+    FuncName = FuncNameRef.str();
+    std::replace(FuncName.begin(), FuncName.end(), '_', '.');
+  }
+  Function *F = M->getFunction(FuncName);
+  if (!F)
+    F = Function::Create(FT, Linkage, FuncName, M);
+  F = cast<Function>(mapValue(BF, F));
   mapFunction(BF, F);
 
+  if (F->isIntrinsic())
+    return F;
+
+  F->setCallingConv(IsKernel ? CallingConv::SPIR_KERNEL
+                             : CallingConv::SPIR_FUNC);
   if (BF->hasDecorate(DecorationReferencedIndirectlyINTEL))
     F->addFnAttr("referenced-indirectly");
-
-  if (!F->isIntrinsic()) {
-    F->setCallingConv(IsKernel ? CallingConv::SPIR_KERNEL
-                               : CallingConv::SPIR_FUNC);
-    if (isFuncNoUnwind())
-      F->addFnAttr(Attribute::NoUnwind);
-    foreachFuncCtlMask(BF,
-                       [&](Attribute::AttrKind Attr) { F->addFnAttr(Attr); });
-  }
+  if (isFuncNoUnwind())
+    F->addFnAttr(Attribute::NoUnwind);
+  foreachFuncCtlMask(BF, [&](Attribute::AttrKind Attr) { F->addFnAttr(Attr); });
 
   for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
        ++I) {

llvm-spirv/lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -46,6 +46,7 @@
 #include "llvm/IR/Operator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/LowerMemIntrinsics.h" // expandMemSetAsLoop()
 
 #include <set>
 #include <vector>
@@ -75,6 +76,15 @@ class SPIRVRegularizeLLVM : public ModulePass {
   void lowerFuncPtr(Function *F, Op OC);
   void lowerFuncPtr(Module *M);
 
+  /// There is no SPIR-V counterpart for @llvm.memset.* intrinsic. Cases with
+  /// constant value and length arguments are emulated via "storing" a constant
+  /// array to the destination. For other cases we wrap the intrinsic in
+  /// @spirv.llvm_memset_* function and expand the intrinsic to a loop via
+  /// expandMemSetAsLoop() from llvm/Transforms/Utils/LowerMemIntrinsics.h
+  /// During reverse translation from SPIR-V to LLVM IR we can detect
+  /// @spirv.llvm_memset_* and replace it with @llvm.memset.
+  void lowerMemset(MemSetInst *MSI);
+
   static char ID;
 
 private:
@@ -84,6 +94,49 @@ class SPIRVRegularizeLLVM : public ModulePass {
 
 char SPIRVRegularizeLLVM::ID = 0;
 
+void SPIRVRegularizeLLVM::lowerMemset(MemSetInst *MSI) {
+  if (isa<Constant>(MSI->getValue()) && isa<ConstantInt>(MSI->getLength()))
+    return; // To be handled in LLVMToSPIRV::transIntrinsicInst
+  Function *IntrinsicFunc = MSI->getCalledFunction();
+  assert(IntrinsicFunc && "Missing function");
+  std::string FuncName = IntrinsicFunc->getName().str();
+  std::replace(FuncName.begin(), FuncName.end(), '.', '_');
+  FuncName = "spirv." + FuncName;
+  if (MSI->isVolatile())
+    FuncName += ".volatile";
+
+  // Redirect @llvm.memset.* call to @spirv.llvm_memset_*
+  Function *F = M->getFunction(FuncName);
+  if (F) {
+    // This function is already linked in.
+    MSI->setCalledFunction(F);
+    return;
+  }
+  // TODO copy arguments attributes: nocapture writeonly.
+  FunctionCallee FC = M->getOrInsertFunction(FuncName, MSI->getFunctionType());
+  MSI->setCalledFunction(FC);
+
+  F = dyn_cast<Function>(FC.getCallee());
+  assert(F && "must be a function!");
+  Argument *Dest = F->getArg(0);
+  Argument *Val = F->getArg(1);
+  Argument *Len = F->getArg(2);
+  Argument *IsVolatile = F->getArg(3);
+  Dest->setName("dest");
+  Val->setName("val");
+  Len->setName("len");
+  IsVolatile->setName("isvolatile");
+  IsVolatile->addAttr(Attribute::ImmArg);
+  BasicBlock *EntryBB = BasicBlock::Create(M->getContext(), "entry", F);
+  IRBuilder<> IRB(EntryBB);
+  auto *MemSet =
+      IRB.CreateMemSet(Dest, Val, Len, MSI->getDestAlign(), MSI->isVolatile());
+  IRB.CreateRetVoid();
+  expandMemSetAsLoop(cast<MemSetInst>(MemSet));
+  MemSet->eraseFromParent();
+  return;
+}
+
 bool SPIRVRegularizeLLVM::runOnModule(Module &Module) {
   M = &Module;
   Ctx = &M->getContext();
@@ -115,8 +168,11 @@ bool SPIRVRegularizeLLVM::regularize() {
         if (auto Call = dyn_cast<CallInst>(&II)) {
           Call->setTailCall(false);
           Function *CF = Call->getCalledFunction();
-          if (CF && CF->isIntrinsic())
+          if (CF && CF->isIntrinsic()) {
             removeFnAttr(Call, Attribute::NoUnwind);
+            if (auto *MSI = dyn_cast<MemSetInst>(Call))
+              lowerMemset(MSI);
+          }
         }
 
         // Remove optimization info not supported by SPIRV

llvm-spirv/test/transcoding/llvm.memset.ll

@@ -6,6 +6,7 @@
 ; RUN: spirv-val %t.spv
 ; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
 
+; CHECK-SPIRV: Decorate [[#NonConstMemset:]] LinkageAttributes "spirv.llvm_memset_p3i8_i32"
 ; CHECK-SPIRV: TypeInt [[Int8:[0-9]+]] 8 0
 ; CHECK-SPIRV: Constant {{[0-9]+}} [[Lenmemset21:[0-9]+]] 4
 ; CHECK-SPIRV: Constant {{[0-9]+}} [[Lenmemset0:[0-9]+]] 12
@@ -19,6 +20,7 @@
 ; CHECK-SPIRV: Variable {{[0-9]+}} [[Val:[0-9]+]] 0 [[Init]]
 ; CHECK-SPIRV: 7 ConstantComposite [[Int8x4]] [[InitComp:[0-9]+]] [[Const21]] [[Const21]] [[Const21]] [[Const21]]
 ; CHECK-SPIRV: Variable {{[0-9]+}} [[ValComp:[0-9]+]] 0 [[InitComp]]
+; CHECK-SPIRV: ConstantFalse [[#]] [[#False:]]
 
 ; CHECK-SPIRV: Bitcast [[Int8Ptr]] [[Target:[0-9]+]] {{[0-9]+}}
 ; CHECK-SPIRV: Bitcast [[Int8PtrConst]] [[Source:[0-9]+]] [[Val]]
@@ -27,6 +29,31 @@
 ; CHECK-SPIRV: Bitcast [[Int8PtrConst]] [[SourceComp:[0-9]+]] [[ValComp]]
 ; CHECK-SPIRV: CopyMemorySized {{[0-9]+}} [[SourceComp]] [[Lenmemset21]] 2 4
 
+; CHECK-SPIRV: FunctionCall [[#]] [[#]] [[#NonConstMemset]] [[#]] [[#]] [[#]] [[#False]]
+
+; CHECK-SPIRV: Function [[#]] [[#NonConstMemset]]
+; CHECK-SPIRV: FunctionParameter [[#]] [[#Dest:]]
+; CHECK-SPIRV: FunctionParameter [[#]] [[#Value:]]
+; CHECK-SPIRV: FunctionParameter [[#]] [[#Len:]]
+; CHECK-SPIRV: FunctionParameter [[#]] [[#Volatile:]]
+
+; CHECK-SPIRV: Label [[#Entry:]]
+; CHECK-SPIRV: IEqual [[#]] [[#IsZeroLen:]] [[#Zero:]] [[#Len]]
+; CHECK-SPIRV: BranchConditional [[#IsZeroLen]] [[#End:]] [[#WhileBody:]]
+
+; CHECK-SPIRV: Label [[#WhileBody]]
+; CHECK-SPIRV: Phi [[#]] [[#Offset:]] [[#Zero]] [[#Entry]] [[#OffsetInc:]] [[#WhileBody]]
+; CHECK-SPIRV: InBoundsPtrAccessChain [[#]] [[#Ptr:]] [[#Dest]] [[#Offset]]
+; CHECK-SPIRV: Store [[#Ptr]] [[#Value]] 2 1
+; CHECK-SPIRV: IAdd [[#]] [[#OffsetInc]] [[#Offset]] [[#One:]]
+; CHECK-SPIRV: ULessThan [[#]] [[#NotEnd:]] [[#OffsetInc]] [[#Len]]
+; CHECK-SPIRV: BranchConditional [[#NotEnd]] [[#WhileBody]] [[#End]]
+
+; CHECK-SPIRV: Label [[#End]]
+; CHECK-SPIRV: Return
+
+; CHECK-SPIRV: FunctionEnd
+
 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-n8:16:32:64"
 target triple = "spir"
 
@@ -36,14 +63,34 @@ target triple = "spir"
 ; CHECK-LLVM: internal unnamed_addr addrspace(2) constant [4 x i8] c"\15\15\15\15"
 
 ; Function Attrs: nounwind
-define spir_func void @_Z5foo11v(%struct.S1 addrspace(4)* noalias nocapture sret %agg.result) #0 {
+define spir_func void @_Z5foo11v(%struct.S1 addrspace(4)* noalias nocapture sret %agg.result, i32 %s1, i64 %s2, i8 %v) #0 {
   %x = alloca [4 x i8]
   %x.bc = bitcast [4 x i8]* %x to i8*
   %1 = bitcast %struct.S1 addrspace(4)* %agg.result to i8 addrspace(4)*
   tail call void @llvm.memset.p4i8.i32(i8 addrspace(4)* align 4 %1, i8 0, i32 12, i1 false)
 ; CHECK-LLVM: call void @llvm.memset.p4i8.i32(i8 addrspace(4)* align 4 %1, i8 0, i32 12, i1 false)
   tail call void @llvm.memset.p0i8.i32(i8* align 4 %x.bc, i8 21, i32 4, i1 false)
 ; CHECK-LLVM: call void @llvm.memcpy.p0i8.p2i8.i32(i8* align 4 %x.bc, i8 addrspace(2)* align 4 %3, i32 4, i1 false)
+
+  ; non-const value
+  tail call void @llvm.memset.p0i8.i32(i8* align 4 %x.bc, i8 %v, i32 3, i1 false)
+; CHECK-LLVM: call void @llvm.memset.p0i8.i32(i8* %x.bc, i8 %v, i32 3, i1 false)
+
+  ; non-const value and size
+  tail call void @llvm.memset.p0i8.i32(i8*  align 4 %x.bc, i8 %v, i32 %s1, i1 false)
+; CHECK-LLVM: call void @llvm.memset.p0i8.i32(i8* %x.bc, i8 %v, i32 %s1, i1 false)
+
+  ; Address spaces, non-const value and size
+  %a = addrspacecast i8 addrspace(4)* %1 to i8 addrspace(3)*
+  tail call void @llvm.memset.p3i8.i32(i8 addrspace(3)* align 4 %a, i8 %v, i32 %s1, i1 false)
+; CHECK-LLVM: call void @llvm.memset.p3i8.i32(i8 addrspace(3)* %a, i8 %v, i32 %s1, i1 false)
+  %b = addrspacecast i8 addrspace(4)* %1 to i8 addrspace(1)*
+  tail call void @llvm.memset.p1i8.i64(i8 addrspace(1)* align 4 %b, i8 %v, i64 %s2, i1 false)
+; CHECK-LLVM: call void @llvm.memset.p1i8.i64(i8 addrspace(1)* %b, i8 %v, i64 %s2, i1 false)
+
+  ; Volatile
+  tail call void @llvm.memset.p1i8.i64(i8 addrspace(1)* align 4 %b, i8 %v, i64 %s2, i1 true)
+; CHECK-LLVM: call void @llvm.memset.p1i8.i64(i8 addrspace(1)* %b, i8 %v, i64 %s2, i1 true)
   ret void
 }
 
@@ -53,6 +100,12 @@ declare void @llvm.memset.p4i8.i32(i8 addrspace(4)* nocapture, i8, i32, i1) #1
 ; Function Attrs: nounwind
 declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i1) #1
 
+; Function Attrs: nounwind
+declare void @llvm.memset.p3i8.i32(i8 addrspace(3)*, i8, i32, i1) #1
+
+; Function Attrs: nounwind
+declare void @llvm.memset.p1i8.i64(i8 addrspace(1)*, i8, i64, i1) #1
+
 attributes #0 = { nounwind "less-precise-fpmad"="false" "no-frame-pointer-elim"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-realign-stack" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
 attributes #1 = { nounwind }