SPV_KHR_untyped_pointers - implement OpUntypedPrefetchKHR (#2752)

When the extension is enabled, we should replace prefetch OpenCL ExtInst with the OpUntypedPrefetchKHR instruction.

Spec:
https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/KHR/SPV_KHR_untyped_pointers.html

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

Author: vmaksimo

llvm-spirv/lib/SPIRV/SPIRVReader.cpp

@@ -2477,6 +2477,36 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
         BV, Builder.CreateIntrinsic(Intrinsic::expect, RetTy, {Val, ExpVal}));
   }
 
+  case OpUntypedPrefetchKHR: {
+    // Do the same as transOCLBuiltinFromExtInst() but for OpUntypedPrefetchKHR.
+    auto *BC = static_cast<SPIRVUntypedPrefetchKHR *>(BV);
+
+    std::vector<Type *> ArgTypes =
+        transTypeVector(BC->getValueTypes(BC->getArguments()), true);
+    Type *RetTy = Type::getVoidTy(*Context);
+
+    std::string MangledName =
+        getSPIRVFriendlyIRFunctionName(OpenCLLIB::Prefetch, ArgTypes, RetTy);
+    opaquifyTypedPointers(ArgTypes);
+
+    FunctionType *FT = FunctionType::get(RetTy, ArgTypes, false);
+    Function *F = M->getFunction(MangledName);
+    if (!F) {
+      F = Function::Create(FT, GlobalValue::ExternalLinkage, MangledName, M);
+      F->setCallingConv(CallingConv::SPIR_FUNC);
+      if (isFuncNoUnwind())
+        F->addFnAttr(Attribute::NoUnwind);
+      if (isFuncReadNone(OCLExtOpMap::map(OpenCLLIB::Prefetch)))
+        F->setDoesNotAccessMemory();
+    }
+
+    auto Args = transValue(BC->getValues(BC->getArguments()), F, BB);
+    CallInst *CI = CallInst::Create(F, Args, BC->getName(), BB);
+    setCallingConv(CI);
+    addFnAttr(CI, Attribute::NoUnwind);
+
+    return mapValue(BV, CI);
+  }
   case OpExtInst: {
     auto *ExtInst = static_cast<SPIRVExtInst *>(BV);
     switch (ExtInst->getExtSetKind()) {

llvm-spirv/lib/SPIRV/SPIRVWriter.cpp

@@ -5336,6 +5336,12 @@ SPIRVValue *LLVMToSPIRVBase::transDirectCallInst(CallInst *CI,
         BM->addExtension(
             ExtensionID::SPV_EXT_relaxed_printf_string_address_space);
       }
+    } else if (DemangledName.find("__spirv_ocl_prefetch") != StringRef::npos) {
+      if (BM->isAllowedToUseExtension(ExtensionID::SPV_KHR_untyped_pointers)) {
+        return BM->addUntypedPrefetchKHRInst(
+            transScavengedType(CI),
+            BM->getIds(transValue(getArguments(CI), BB)), BB);
+      }
     }
 
     return addDecorations(

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -4200,5 +4200,101 @@ _SPIRV_OP(ConvertHandleToSamplerINTEL)
 _SPIRV_OP(ConvertHandleToSampledImageINTEL)
 #undef _SPIRV_OP
 
+class SPIRVUntypedPrefetchKHR : public SPIRVInstruction {
+public:
+  static const Op OC = OpUntypedPrefetchKHR;
+  static const SPIRVWord FixedWordCount = 3;
+
+  SPIRVUntypedPrefetchKHR(SPIRVType *Ty, std::vector<SPIRVWord> &TheArgs,
+                          SPIRVBasicBlock *BB)
+      : SPIRVInstruction(FixedWordCount, OC, BB) {
+    setHasNoId();
+    setHasNoType();
+    PtrTy = TheArgs[0];
+    NumBytes = TheArgs[1];
+    if (TheArgs.size() > 2)
+      RW.push_back(TheArgs[2]);
+    if (TheArgs.size() > 3)
+      Locality.push_back(TheArgs[3]);
+    if (TheArgs.size() > 4)
+      CacheTy.push_back(TheArgs[4]);
+    assert(BB && "Invalid BB");
+    validate();
+  }
+
+  SPIRVUntypedPrefetchKHR() : SPIRVInstruction(OC) {
+    setHasNoId();
+    setHasNoType();
+  }
+
+  void validate() const override {
+    SPIRVInstruction::validate();
+    SPIRVErrorLog &SPVErrLog = this->getModule()->getErrorLog();
+    std::string InstName = "OpUntypedPrefetchKHR";
+    SPVErrLog.checkError(getValueType(PtrTy)->isTypePointer(),
+                         SPIRVEC_InvalidInstruction,
+                         InstName + "\nFirst argument must be a pointer\n");
+    SPVErrLog.checkError(
+        getValueType(PtrTy)->getPointerStorageClass() ==
+            StorageClassCrossWorkgroup,
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nFirst argument must be a pointer in CrossWorkgroup "
+                   "storage class\n");
+    SPVErrLog.checkError(
+        getValueType(NumBytes)->isTypeInt(), SPIRVEC_InvalidInstruction,
+        InstName + "\nSecond argument (Num Bytes) must be an integer\n");
+    SPVErrLog.checkError(
+        RW.empty() || (RW.size() == 1 && getValueType(RW[0])->isTypeInt()),
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nThird argument (RW) must be an integer\n");
+    SPVErrLog.checkError(
+        Locality.empty() ||
+            (Locality.size() == 1 && getValueType(Locality[0])->isTypeInt()),
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nFourth argument (Locality) must be an integer\n");
+    SPVErrLog.checkError(
+        CacheTy.empty() ||
+            (CacheTy.size() == 1 && getValueType(CacheTy[0])->isTypeInt()),
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nFifth argument (Cache Type) must be an integer\n");
+  }
+
+  void setWordCount(SPIRVWord TheWordCount) override {
+    SPIRVEntry::setWordCount(TheWordCount);
+    if (TheWordCount > 3)
+      RW.resize(1);
+    if (TheWordCount > 4)
+      Locality.resize(1);
+    if (TheWordCount > 5)
+      CacheTy.resize(1);
+  }
+  const std::vector<SPIRVWord> getArguments() const {
+    std::vector<SPIRVWord> Args;
+    Args.push_back(PtrTy);
+    Args.push_back(NumBytes);
+    if (!RW.empty())
+      Args.push_back(RW[0]);
+    if (!Locality.empty())
+      Args.push_back(Locality[0]);
+    if (!CacheTy.empty())
+      Args.push_back(CacheTy[0]);
+    return Args;
+  }
+
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(CapabilityUntypedPointersKHR);
+  }
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_KHR_untyped_pointers;
+  }
+  _SPIRV_DEF_ENCDEC5(PtrTy, NumBytes, RW, Locality, CacheTy)
+protected:
+  SPIRVId PtrTy;
+  SPIRVId NumBytes;
+  std::vector<SPIRVId> RW;
+  std::vector<SPIRVId> Locality;
+  std::vector<SPIRVId> CacheTy;
+};
+
 } // namespace SPIRV
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -493,6 +493,9 @@ class SPIRVModuleImpl : public SPIRVModule {
   SPIRVInstruction *addExpectKHRInst(SPIRVType *ResultTy, SPIRVValue *Value,
                                      SPIRVValue *ExpectedValue,
                                      SPIRVBasicBlock *BB) override;
+  SPIRVInstruction *addUntypedPrefetchKHRInst(SPIRVType *Ty,
+                                              std::vector<SPIRVWord> Args,
+                                              SPIRVBasicBlock *BB) override;
 
   virtual SPIRVId getExtInstSetId(SPIRVExtInstSetKind Kind) const override;
 
@@ -1832,6 +1835,11 @@ SPIRVInstruction *SPIRVModuleImpl::addExpectKHRInst(SPIRVType *ResultTy,
                         BB);
 }
 
+SPIRVInstruction *SPIRVModuleImpl::addUntypedPrefetchKHRInst(
+    SPIRVType *Ty, std::vector<SPIRVWord> Args, SPIRVBasicBlock *BB) {
+  return addInstruction(new SPIRVUntypedPrefetchKHR(Ty, Args, BB), BB);
+}
+
 // Create AliasDomainDeclINTEL/AliasScopeDeclINTEL/AliasScopeListDeclINTEL
 // instructions
 template <typename AliasingInstType>

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -500,6 +500,9 @@ class SPIRVModule {
                                              SPIRVValue *Value,
                                              SPIRVValue *ExpectedValue,
                                              SPIRVBasicBlock *BB) = 0;
+  virtual SPIRVInstruction *
+  addUntypedPrefetchKHRInst(SPIRVType *Ty, std::vector<SPIRVWord> Args,
+                            SPIRVBasicBlock *BB) = 0;
 
   virtual SPIRVId getExtInstSetId(SPIRVExtInstSetKind Kind) const = 0;
 

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVOpCodeEnum.h

@@ -339,6 +339,7 @@ _SPIRV_OP(UntypedAccessChainKHR, 4419)
 _SPIRV_OP(UntypedInBoundsAccessChainKHR, 4420)
 _SPIRV_OP(UntypedPtrAccessChainKHR, 4423)
 _SPIRV_OP(UntypedInBoundsPtrAccessChainKHR, 4424)
+_SPIRV_OP(UntypedPrefetchKHR, 4426)
 _SPIRV_OP(GroupNonUniformRotateKHR, 4431)
 _SPIRV_OP(SDotKHR, 4450)
 _SPIRV_OP(UDotKHR, 4451)

llvm-spirv/spirv-headers-tag.conf

@@ -1 +1 @@
-db5a00f8cebe81146cafabf89019674a3c4bf03d
+efb6b4099ddb8fa60f62956dee592c4b94ec6a49

llvm-spirv/test/extensions/INTEL/SPV_INTEL_cache_controls/decorate-prefetch-w-cache-controls.ll

@@ -1,7 +1,14 @@
 ; RUN: llvm-as %s -o %t.bc
-; RUN: llvm-spirv --spirv-ext=+SPV_INTEL_cache_controls -spirv-text %t.bc -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv --spirv-ext=+SPV_INTEL_cache_controls -spirv-text %t.bc -o - | FileCheck %s --check-prefixes=CHECK-SPIRV,CHECK-SPIRV-TYPED-PTRS
 
 ; RUN: llvm-spirv --spirv-ext=+SPV_INTEL_cache_controls %t.bc -o %t.spv
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -r %t.spv --spirv-target-env=SPV-IR -o - | llvm-dis -o - | FileCheck %s --check-prefix=CHECK-LLVM
+
+; RUN: llvm-spirv --spirv-ext=+SPV_INTEL_cache_controls,+SPV_KHR_untyped_pointers -spirv-text %t.bc -o - | FileCheck %s --check-prefixes=CHECK-SPIRV,CHECK-SPIRV-UNTYPED-PTRS
+
+; RUN: llvm-spirv --spirv-ext=+SPV_INTEL_cache_controls,+SPV_KHR_untyped_pointers %t.bc -o %t.spv
+; RUN: spirv-val %t.spv
 ; RUN: llvm-spirv -r %t.spv --spirv-target-env=SPV-IR -o - | llvm-dis -o - | FileCheck %s --check-prefix=CHECK-LLVM
 
 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"
@@ -24,9 +31,12 @@ $_ZTSZ4mainEUlvE_ = comdat any
 ; CHECK-SPIRV: Decorate [[PTR_ID2:.*]] CacheControlLoadINTEL 1 1
 ; CHECK-SPIRV: Decorate [[PTR_ID3:.*]] CacheControlLoadINTEL 2 3
 
-; CHECK-SPIRV: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID1]] [[#]]
-; CHECK-SPIRV: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID2]] [[#]]
-; CHECK-SPIRV: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID3]] [[#]]
+; CHECK-SPIRV-TYPED-PTRS: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID1]] [[#]]
+; CHECK-SPIRV-TYPED-PTRS: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID2]] [[#]]
+; CHECK-SPIRV-TYPED-PTRS: ExtInst [[#]] [[#]] [[#]] prefetch [[PTR_ID3]] [[#]]
+; CHECK-SPIRV-UNTYPED-PTRS: UntypedPrefetchKHR [[PTR_ID1]] [[#]]
+; CHECK-SPIRV-UNTYPED-PTRS: UntypedPrefetchKHR [[PTR_ID2]] [[#]]
+; CHECK-SPIRV-UNTYPED-PTRS: UntypedPrefetchKHR [[PTR_ID3]] [[#]]
 
 ; Check that the appropriate !spirv.Decorations are preserved after reverse
 ; translation