[spirv] Add Vk::RawBufferLoad intrinsic (#4069)

yuriy-odonnell-epic · web-flow · commit f83655590783 · 2021-11-17T15:41:54.000-05:00
* Only support raw buffer (i.e. HLSL ByteAddressBuffer) style loads of single uint-s
* Always use 4 byte alignment for loads (same as HLSL ByteAddressBuffer)
* Add PhysicalStorageBufferAddresses capability and KHR_physical_storage_buffer extension when needed
* Promote memory addressing model to PhysicalStorageBuffer64 when needed
* Add custom alignment support to SpirvLoad
* Add createUnaryOp() overload that takes SpirvType
* Add getPhysicalStorageBufferType()
* Add documentation and a test case for the new intrinsic
diff --git a/docs/SPIR-V.rst b/docs/SPIR-V.rst
@@ -3714,6 +3714,7 @@ implicit ``vk`` namepsace.
     const uint QueueFamilyScope = 5;
   
     uint64_t ReadClock(in uint scope);
+    uint     RawBufferLoad(in uint64_t deviceAddress);
   } // end namespace
 
 
@@ -3751,6 +3752,35 @@ For example:
 
   uint64_t clock = vk::ReadClock(vk::SubgroupScope);
 
+RawBufferLoad
+~~~~~~~~~~~~~
+This intrinsic funcion has the following signature:
+
+.. code:: hlsl
+
+  uint RawBufferLoad(in uint64_t deviceAddress);
+
+This exposes a subset of the `VK_KHR_buffer_device_address <https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_buffer_device_address.html>`_
+and `SPV_KHR_physical_storage_buffer <https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/KHR/SPV_KHR_physical_storage_buffer.asciidoc>`_ 
+functionality to HLSL. 
+
+It allows the shader program to load a single 32 bit value from a GPU
+accessible memory at given address, similar to ``ByteAddressBuffer.Load()``.
+Like ``ByteAddressBuffer``, this intrinsic requires a 4 byte aligned address.
+
+Using this intrinsic adds ``PhysicalStorageBufferAddresses`` capability and 
+``SPV_KHR_physical_storage_buffer`` extension requirements as well as changing 
+the addressing model to ``PhysicalStorageBuffer64``.
+
+Example:
+
+.. code:: hlsl
+
+  uint64_t Address;
+  float4 main() : SV_Target0 {
+    uint Value = vk::RawBufferLoad(Address);
+    return asfloat(Value);
+  }
 
 Supported Command-line Options
 ==============================
diff --git a/include/dxc/HlslIntrinsicOp.h b/include/dxc/HlslIntrinsicOp.h
@@ -221,6 +221,9 @@ enum class IntrinsicOp {  IOP_AcceptHitAndEndSearch,
   IOP_unpack_s8s32,
   IOP_unpack_u8u16,
   IOP_unpack_u8u32,
+#ifdef ENABLE_SPIRV_CODEGEN
+  IOP_VkRawBufferLoad,
+#endif // ENABLE_SPIRV_CODEGEN
 #ifdef ENABLE_SPIRV_CODEGEN
   IOP_VkReadClock,
 #endif // ENABLE_SPIRV_CODEGEN
diff --git a/lib/HLSL/HLOperationLower.cpp b/lib/HLSL/HLOperationLower.cpp
@@ -5661,6 +5661,7 @@ IntrinsicLower gLowerTable[] = {
     {IntrinsicOp::IOP_unpack_u8u32, TranslateUnpack, DXIL::OpCode::Unpack4x8},
 #ifdef ENABLE_SPIRV_CODEGEN
     { IntrinsicOp::IOP_VkReadClock, UnsupportedVulkanIntrinsic, DXIL::OpCode::NumOpCodes },
+    { IntrinsicOp::IOP_VkRawBufferLoad, UnsupportedVulkanIntrinsic, DXIL::OpCode::NumOpCodes },
 #endif // ENABLE_SPIRV_CODEGEN
     {IntrinsicOp::MOP_Append, StreamOutputLower, DXIL::OpCode::EmitStream},
     {IntrinsicOp::MOP_RestartStrip, StreamOutputLower, DXIL::OpCode::CutStream},
diff --git a/tools/clang/include/clang/SPIRV/FeatureManager.h b/tools/clang/include/clang/SPIRV/FeatureManager.h
@@ -52,6 +52,7 @@ enum class Extension {
   NV_mesh_shader,
   KHR_ray_query,
   EXT_shader_image_int64,
+  KHR_physical_storage_buffer,
   Unknown,
 };
 
diff --git a/tools/clang/include/clang/SPIRV/SpirvBuilder.h b/tools/clang/include/clang/SPIRV/SpirvBuilder.h
@@ -212,6 +212,8 @@ class SpirvBuilder {
   /// the instruction pointer for the result.
   SpirvUnaryOp *createUnaryOp(spv::Op op, QualType resultType,
                               SpirvInstruction *operand, SourceLocation loc);
+  SpirvUnaryOp *createUnaryOp(spv::Op op, const SpirvType *resultType,
+                              SpirvInstruction *operand, SourceLocation loc);
 
   /// \brief Creates a binary operation with the given SPIR-V opcode. Returns
   /// the instruction pointer for the result.
@@ -689,6 +691,10 @@ class SpirvBuilder {
 
   SpirvString *getString(llvm::StringRef str);
 
+  const HybridPointerType *getPhysicalStorageBufferType(QualType pointee);
+  const SpirvPointerType *
+  getPhysicalStorageBufferType(const SpirvType *pointee);
+
 public:
   std::vector<uint32_t> takeModule();
 
@@ -793,8 +799,14 @@ class SpirvBuilder {
 
 void SpirvBuilder::requireCapability(spv::Capability cap, SourceLocation loc) {
   auto *capability = new (context) SpirvCapability(loc, cap);
-  if (!mod->addCapability(capability))
+  if (mod->addCapability(capability)) {
+    if (cap == spv::Capability::PhysicalStorageBufferAddresses) {
+      mod->promoteAddressingModel(
+          spv::AddressingModel::PhysicalStorageBuffer64);
+    }
+  } else {
     capability->releaseMemory();
+  }
 }
 
 void SpirvBuilder::requireExtension(llvm::StringRef ext, SourceLocation loc) {
diff --git a/tools/clang/include/clang/SPIRV/SpirvInstruction.h b/tools/clang/include/clang/SPIRV/SpirvInstruction.h
@@ -332,6 +332,9 @@ class SpirvMemoryModel : public SpirvInstruction {
 
   spv::AddressingModel getAddressingModel() const { return addressModel; }
   spv::MemoryModel getMemoryModel() const { return memoryModel; }
+  void setAddressingModel(spv::AddressingModel addrModel) {
+    addressModel = addrModel;
+  }
 
 private:
   spv::AddressingModel addressModel;
@@ -1669,9 +1672,14 @@ class SpirvLoad : public SpirvInstruction {
     return memoryAccess.getValue();
   }
 
+  void setAlignment(uint32_t alignment);
+  bool hasAlignment() const { return memoryAlignment.hasValue(); }
+  uint32_t getAlignment() const { return memoryAlignment.getValue(); }
+
 private:
   SpirvInstruction *pointer;
   llvm::Optional<spv::MemoryAccessMask> memoryAccess;
+  llvm::Optional<uint32_t> memoryAlignment;
 };
 
 /// \brief OpCopyObject instruction
@@ -1869,6 +1877,9 @@ class SpirvUnaryOp : public SpirvInstruction {
   SpirvUnaryOp(spv::Op opcode, QualType resultType, SourceLocation loc,
                SpirvInstruction *op);
 
+  SpirvUnaryOp(spv::Op opcode, const SpirvType *resultType, SourceLocation loc,
+               SpirvInstruction *op);
+
   DEFINE_RELEASE_MEMORY_FOR_CLASS(SpirvUnaryOp)
 
   // For LLVM-style RTTI
diff --git a/tools/clang/include/clang/SPIRV/SpirvModule.h b/tools/clang/include/clang/SPIRV/SpirvModule.h
@@ -105,6 +105,12 @@ class SpirvModule {
   // Set the memory model of the module.
   void setMemoryModel(SpirvMemoryModel *model);
 
+  // Increases addressing model requirement for the module:
+  // Logical -> Physical32 -> Physical64 -> PhysicalStorageBuffer64.
+  // Requires setMemoryModel() to be called first to set the base memory model.
+  // Returns true if addressing model was changed.
+  bool promoteAddressingModel(spv::AddressingModel addrModel);
+
   // Add an entry point to the module.
   void addEntryPoint(SpirvEntryPoint *);
 
diff --git a/tools/clang/lib/SPIRV/CapabilityVisitor.cpp b/tools/clang/lib/SPIRV/CapabilityVisitor.cpp
@@ -189,6 +189,11 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
   // Pointer type
   else if (const auto *ptrType = dyn_cast<SpirvPointerType>(type)) {
     addCapabilityForType(ptrType->getPointeeType(), loc, sc);
+    if (sc == spv::StorageClass::PhysicalStorageBuffer) {
+      addExtension(Extension::KHR_physical_storage_buffer,
+                   "SPV_KHR_physical_storage_buffer", loc);
+      addCapability(spv::Capability::PhysicalStorageBufferAddresses);
+    }
   }
   // Struct type
   else if (const auto *structType = dyn_cast<StructType>(type)) {
@@ -525,8 +530,7 @@ bool CapabilityVisitor::visitInstruction(SpirvInstruction *instr) {
   case spv::Op::OpRayQueryInitializeKHR: {
     auto rayQueryInst = dyn_cast<SpirvRayQueryOpKHR>(instr);
     if (rayQueryInst->hasCullFlags()) {
-      addCapability(
-          spv::Capability::RayTraversalPrimitiveCullingKHR);
+      addCapability(spv::Capability::RayTraversalPrimitiveCullingKHR);
     }
 
     break;
@@ -581,6 +585,7 @@ bool CapabilityVisitor::visit(SpirvEntryPoint *entryPoint) {
     llvm_unreachable("found unknown shader model");
     break;
   }
+
   return true;
 }
 
diff --git a/tools/clang/lib/SPIRV/EmitVisitor.cpp b/tools/clang/lib/SPIRV/EmitVisitor.cpp
@@ -1201,8 +1201,15 @@ bool EmitVisitor::visit(SpirvLoad *inst) {
   curInst.push_back(inst->getResultTypeId());
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst));
   curInst.push_back(getOrAssignResultId<SpirvInstruction>(inst->getPointer()));
-  if (inst->hasMemoryAccessSemantics())
-    curInst.push_back(static_cast<uint32_t>(inst->getMemoryAccess()));
+  if (inst->hasMemoryAccessSemantics()) {
+    spv::MemoryAccessMask memoryAccess = inst->getMemoryAccess();
+    curInst.push_back(static_cast<uint32_t>(memoryAccess));
+    if (inst->hasAlignment()) {
+      assert(static_cast<uint32_t>(memoryAccess) &
+             static_cast<uint32_t>(spv::MemoryAccessMask::Aligned));
+      curInst.push_back(inst->getAlignment());
+    }
+  }
   finalizeInstruction(&mainBinary);
   emitDebugNameForInstruction(getOrAssignResultId<SpirvInstruction>(inst),
                               inst->getDebugName());
diff --git a/tools/clang/lib/SPIRV/FeatureManager.cpp b/tools/clang/lib/SPIRV/FeatureManager.cpp
@@ -143,6 +143,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_KHR_fragment_shading_rate",
             Extension::KHR_fragment_shading_rate)
       .Case("SPV_EXT_shader_image_int64", Extension::EXT_shader_image_int64)
+      .Case("SPV_KHR_physical_storage_buffer",
+            Extension::KHR_physical_storage_buffer)
       .Default(Extension::Unknown);
 }
 
@@ -196,6 +198,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_KHR_fragment_shading_rate";
   case Extension::EXT_shader_image_int64:
     return "SPV_EXT_shader_image_int64";
+  case Extension::KHR_physical_storage_buffer:
+    return "SPV_KHR_physical_storage_buffer";
   default:
     break;
   }
diff --git a/tools/clang/lib/SPIRV/SpirvBuilder.cpp b/tools/clang/lib/SPIRV/SpirvBuilder.cpp
@@ -334,6 +334,16 @@ SpirvUnaryOp *SpirvBuilder::createUnaryOp(spv::Op op, QualType resultType,
   return instruction;
 }
 
+SpirvUnaryOp *SpirvBuilder::createUnaryOp(spv::Op op,
+                                          const SpirvType *resultType,
+                                          SpirvInstruction *operand,
+                                          SourceLocation loc) {
+  assert(insertPoint && "null insert point");
+  auto *instruction = new (context) SpirvUnaryOp(op, resultType, loc, operand);
+  insertPoint->addInstruction(instruction);
+  return instruction;
+}
+
 SpirvBinaryOp *SpirvBuilder::createBinaryOp(spv::Op op, QualType resultType,
                                             SpirvInstruction *lhs,
                                             SpirvInstruction *rhs,
@@ -1530,6 +1540,18 @@ SpirvString *SpirvBuilder::getString(llvm::StringRef str) {
   return instr;
 }
 
+const HybridPointerType *
+SpirvBuilder::getPhysicalStorageBufferType(QualType pointee) {
+  return context.getPointerType(pointee,
+                                spv::StorageClass::PhysicalStorageBuffer);
+}
+
+const SpirvPointerType *
+SpirvBuilder::getPhysicalStorageBufferType(const SpirvType *pointee) {
+  return context.getPointerType(pointee,
+                                spv::StorageClass::PhysicalStorageBuffer);
+}
+
 void SpirvBuilder::addModuleInitCallToEntryPoints() {
   if (moduleInit == nullptr)
     return;
diff --git a/tools/clang/lib/SPIRV/SpirvEmitter.cpp b/tools/clang/lib/SPIRV/SpirvEmitter.cpp
@@ -687,6 +687,7 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
   }
 
   // Addressing and memory model are required in a valid SPIR-V module.
+  // It may be promoted based on features used by this shader.
   spvBuilder.setMemoryModel(spv::AddressingModel::Logical,
                             spv::MemoryModel::GLSL450);
 
@@ -7633,6 +7634,9 @@ SpirvEmitter::processIntrinsicCallExpr(const CallExpr *callExpr) {
   case hlsl::IntrinsicOp::IOP_VkReadClock:
     retVal = processIntrinsicReadClock(callExpr);
     break;
+  case hlsl::IntrinsicOp::IOP_VkRawBufferLoad:
+    retVal = processRawBufferLoad(callExpr);
+    break;
   case hlsl::IntrinsicOp::IOP_saturate:
     retVal = processIntrinsicSaturate(callExpr);
     break;
@@ -12539,6 +12543,32 @@ SpirvEmitter::processSpvIntrinsicCallExpr(const CallExpr *expr) {
   return retVal;
 }
 
+SpirvInstruction *SpirvEmitter::processRawBufferLoad(const CallExpr *callExpr) {
+  clang::SourceLocation loc = callExpr->getExprLoc();
+  const clang::Expr *addressExpr = callExpr->getArg(0);
+  SpirvInstruction *address = doExpr(addressExpr);
+
+  const SpirvPointerType *bufferType =
+      spvBuilder.getPhysicalStorageBufferType(spvContext.getUIntType(32));
+
+  SpirvUnaryOp *bufferReference =
+      spvBuilder.createUnaryOp(spv::Op::OpBitcast, bufferType, address, loc);
+
+  bufferReference->setStorageClass(spv::StorageClass::PhysicalStorageBuffer);
+
+  SpirvAccessChain *ac = spvBuilder.createAccessChain(astContext.UnsignedIntTy,
+                                                      bufferReference, {}, loc);
+
+  SpirvLoad *loadInst =
+      spvBuilder.createLoad(astContext.UnsignedIntTy, ac, loc);
+
+  // Raw buffer loads have the same alignment requirement as
+  // ByteAddressBuffer in HLSL
+  loadInst->setAlignment(4);
+
+  return loadInst;
+}
+
 bool SpirvEmitter::spirvToolsValidate(std::vector<uint32_t> *mod,
                                       std::string *messages) {
   spvtools::SpirvTools tools(featureManager.getTargetEnv());
diff --git a/tools/clang/lib/SPIRV/SpirvEmitter.h b/tools/clang/lib/SPIRV/SpirvEmitter.h
@@ -597,6 +597,9 @@ class SpirvEmitter : public ASTConsumer {
   /// Process spirv intrinsic instruction
   SpirvInstruction *processSpvIntrinsicCallExpr(const CallExpr *expr);
 
+  /// Custom intrinsic to support basic buffer_reference use case
+  SpirvInstruction *processRawBufferLoad(const CallExpr *callExpr);
+
 private:
   /// Returns the <result-id> for constant value 0 of the given type.
   SpirvConstant *getValueZero(QualType type);
diff --git a/tools/clang/lib/SPIRV/SpirvInstruction.cpp b/tools/clang/lib/SPIRV/SpirvInstruction.cpp
@@ -755,6 +755,18 @@ SpirvLoad::SpirvLoad(QualType resultType, SourceLocation loc,
     : SpirvInstruction(IK_Load, spv::Op::OpLoad, resultType, loc, range),
       pointer(pointerInst), memoryAccess(mask) {}
 
+void SpirvLoad::setAlignment(uint32_t alignment) {
+  assert(alignment != 0);
+  assert(llvm::isPowerOf2_32(alignment));
+  if (!memoryAccess.hasValue()) {
+    memoryAccess = spv::MemoryAccessMask::Aligned;
+  } else {
+    memoryAccess.getValue() =
+        memoryAccess.getValue() | spv::MemoryAccessMask::Aligned;
+  }
+  memoryAlignment = alignment;
+}
+
 SpirvCopyObject::SpirvCopyObject(QualType resultType, SourceLocation loc,
                                  SpirvInstruction *pointerInst)
     : SpirvInstruction(IK_CopyObject, spv::Op::OpCopyObject, resultType, loc),
@@ -802,6 +814,12 @@ SpirvUnaryOp::SpirvUnaryOp(spv::Op opcode, QualType resultType,
                            SourceLocation loc, SpirvInstruction *op)
     : SpirvInstruction(IK_UnaryOp, opcode, resultType, loc), operand(op) {}
 
+SpirvUnaryOp::SpirvUnaryOp(spv::Op opcode, const SpirvType *resultType,
+                           SourceLocation loc, SpirvInstruction *op)
+    : SpirvInstruction(IK_UnaryOp, opcode, QualType(), loc), operand(op) {
+  setResultType(resultType);
+}
+
 bool SpirvUnaryOp::isConversionOp() const {
   return opcode == spv::Op::OpConvertFToU || opcode == spv::Op::OpConvertFToS ||
          opcode == spv::Op::OpConvertSToF || opcode == spv::Op::OpConvertUToF ||
diff --git a/tools/clang/lib/SPIRV/SpirvModule.cpp b/tools/clang/lib/SPIRV/SpirvModule.cpp
@@ -241,6 +241,34 @@ void SpirvModule::setMemoryModel(SpirvMemoryModel *model) {
   memoryModel = model;
 }
 
+bool SpirvModule::promoteAddressingModel(spv::AddressingModel addrModel) {
+  assert(memoryModel && "base memory model must be set first");
+  auto getPriority = [](spv::AddressingModel am) -> int {
+    switch (am) {
+    default:
+      assert(false && "unknown addressing model");
+    case spv::AddressingModel::Logical:
+      return 0;
+    case spv::AddressingModel::Physical32:
+      return 1;
+    case spv::AddressingModel::Physical64:
+      return 2;
+    case spv::AddressingModel::PhysicalStorageBuffer64:
+      return 3;
+    }
+  };
+
+  int current = getPriority(memoryModel->getAddressingModel());
+  int pending = getPriority(addrModel);
+
+  if (pending > current) {
+    memoryModel->setAddressingModel(addrModel);
+    return true;
+  } else {
+    return false;
+  }
+}
+
 void SpirvModule::addEntryPoint(SpirvEntryPoint *ep) {
   assert(ep && "cannot add null as an entry point");
   entryPoints.push_back(ep);
diff --git a/tools/clang/test/CodeGenSPIRV/intrinsics.vkrawbufferload.hlsl b/tools/clang/test/CodeGenSPIRV/intrinsics.vkrawbufferload.hlsl
diff --git a/tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp b/tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp
diff --git a/utils/hct/gen_intrin_main.txt b/utils/hct/gen_intrin_main.txt
