[SPIRV] Add write to image buffer for shaders.

llvm/docs/SPIRVUsage.rst

@@ -400,6 +400,11 @@ SPIR-V backend, along with their descriptions and argument details.
        return type is a scalar, then the first element of the vector is \
        returned. If the return type is an n-element vector, then the first \
        n-elements of the 4-element vector are returned.
+   * - `int_spv_typedBufferStore`
+     - void
+     - `[spirv.Image Image, 32-bit Integer coordinate, vec4 data]`
+     - Stores the data to the image buffer at the given coordinate. The \
+       data must be a 4-element vector.
 
 .. _spirv-builtin-functions:
 

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -112,4 +112,10 @@ let TargetPrefix = "spv" in {
   // vector.
   def int_spv_typedBufferLoad
       : DefaultAttrsIntrinsic<[llvm_any_ty], [llvm_any_ty, llvm_i32_ty]>;
+
+  // Write a value to the image buffer. Translates directly to a single
+  // OpImageWrite.
+  def int_spv_typedBufferStore
+    : DefaultAttrsIntrinsic<[], [llvm_any_ty, llvm_i32_ty, llvm_anyvector_ty]>;
+
 }

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -271,6 +271,8 @@ class SPIRVInstructionSelector : public InstructionSelector {
   void selectReadImageIntrinsic(Register &ResVReg, const SPIRVType *ResType,
                                 MachineInstr &I) const;
 
+  void selectImageWriteIntrinsic(MachineInstr &I) const;
+
   // Utilities
   std::pair<Register, bool>
   buildI32Constant(uint32_t Val, MachineInstr &I,
@@ -2853,6 +2855,10 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   case Intrinsic::spv_handle_fromBinding: {
     return selectHandleFromBinding(ResVReg, ResType, I);
   }
+  case Intrinsic::spv_typedBufferStore: {
+    selectImageWriteIntrinsic(I);
+    return true;
+  }
   case Intrinsic::spv_typedBufferLoad: {
     selectReadImageIntrinsic(ResVReg, ResType, I);
     return true;
@@ -2971,6 +2977,23 @@ void SPIRVInstructionSelector::extractSubvector(
     MIB.addUse(ComponentReg);
 }
 
+void SPIRVInstructionSelector::selectImageWriteIntrinsic(
+    MachineInstr &I) const {
+  // If the load of the image is in a different basic block, then
+  // this will generate invalid code. A proper solution is to move
+  // the OpLoad from selectHandleFromBinding here. However, to do
+  // that we will need to change the return type of the intrinsic.
+  // We will do that when we can, but for now trying to move forward with other
+  // issues.
+  Register DataReg = I.getOperand(3).getReg();
+  assert(GR.getResultType(DataReg)->getOpcode() == SPIRV::OpTypeVector);
+  assert(GR.getScalarOrVectorComponentCount(GR.getResultType(DataReg)) == 4);
+  BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(SPIRV::OpImageWrite))
+      .addUse(I.getOperand(1).getReg())
+      .addUse(I.getOperand(2).getReg())
+      .addUse(DataReg);
+}
+
 Register SPIRVInstructionSelector::buildPointerToResource(
     const SPIRVType *ResType, uint32_t Set, uint32_t Binding,
     uint32_t ArraySize, Register IndexReg, bool IsNonUniform,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -661,6 +661,7 @@ void RequirementHandler::initAvailableCapabilitiesForOpenCL(
   addAvailableCaps({Capability::Addresses, Capability::Float16Buffer,
                     Capability::Kernel, Capability::Vector16,
                     Capability::Groups, Capability::GenericPointer,
+                    Capability::StorageImageWriteWithoutFormat,
                     Capability::StorageImageReadWithoutFormat});
   if (ST.hasOpenCLFullProfile())
     addAvailableCaps({Capability::Int64, Capability::Int64Atomics});
@@ -723,7 +724,8 @@ void RequirementHandler::initAvailableCapabilitiesForVulkan(
 
   // Became core in Vulkan 1.3
   if (ST.isAtLeastSPIRVVer(VersionTuple(1, 6)))
-    addAvailableCaps({Capability::StorageImageReadWithoutFormat});
+    addAvailableCaps({Capability::StorageImageWriteWithoutFormat,
+                      Capability::StorageImageReadWithoutFormat});
 }
 
 } // namespace SPIRV
@@ -1430,6 +1432,13 @@ void addInstrRequirements(const MachineInstr &MI,
       Reqs.addCapability(SPIRV::Capability::StorageImageReadWithoutFormat);
     break;
   }
+  case SPIRV::OpImageWrite: {
+    Register ImageReg = MI.getOperand(0).getReg();
+    SPIRVType *TypeDef = ST.getSPIRVGlobalRegistry()->getResultType(ImageReg);
+    if (isImageTypeWithUnknownFormat(TypeDef))
+      Reqs.addCapability(SPIRV::Capability::StorageImageWriteWithoutFormat);
+    break;
+  }
 
   default:
     break;

llvm/test/CodeGen/SPIRV/hlsl-resources/BufferStore.ll

@@ -0,0 +1,37 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv-vulkan-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-vulkan-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK-NOT: OpCapability StorageImageReadWithoutFormat
+
+; CHECK-DAG: OpDecorate [[IntBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[IntBufferVar]] Binding 7
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[zero:%[0-9]+]] = OpConstant [[int]] 0
+; CHECK-DAG: [[v4_int:%[0-9]+]] = OpTypeVector [[int]] 4
+; CHECK-DAG: [[RWBufferTypeInt:%[0-9]+]] = OpTypeImage [[int]] Buffer 2 0 0 2 R32i {{$}}
+; CHECK-DAG: [[IntBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeInt]]
+; CHECK-DAG: [[IntBufferVar]] = OpVariable [[IntBufferPtrType]] UniformConstant
+
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+declare <4 x i32> @get_data() #1
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferStore_Vec4_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_i32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: [[data:%[0-9]+]] = OpFunctionCall
+  %data = call <4 x i32> @get_data()
+; CHECK: OpImageWrite [[buffer]] [[zero]] [[data]]
+  call void @llvm.spv.typedBufferStore(target("spirv.Image", i32, 5, 2, 0, 0, 2, 24) %buffer0, i32 0, <4 x i32> %data)
+
+  ret void
+}
+
+attributes #0 = { convergent noinline norecurse "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+attributes #1 = { convergent noinline norecurse "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }

llvm/test/CodeGen/SPIRV/hlsl-resources/UnknownBufferStore.ll

@@ -0,0 +1,36 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv1.6-vulkan1.3-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-vulkan1.3-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpCapability StorageImageWriteWithoutFormat
+; CHECK-DAG: OpDecorate [[IntBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[IntBufferVar]] Binding 7
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[ten:%[0-9]+]] = OpConstant [[int]] 10
+; CHECK-DAG: [[v4_int:%[0-9]+]] = OpTypeVector [[int]] 4
+; CHECK-DAG: [[RWBufferTypeInt:%[0-9]+]] = OpTypeImage [[int]] Buffer 2 0 0 2 Unknown {{$}}
+; CHECK-DAG: [[IntBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeInt]]
+; CHECK-DAG: [[IntBufferVar]] = OpVariable [[IntBufferPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+declare <4 x i32> @get_data() #1
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad_Vec4_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 0)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_0(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: [[data:%[0-9]+]] = OpFunctionCall
+  %data = call <4 x i32> @get_data()
+; CHECK: OpImageWrite [[buffer]] [[ten]] [[data]]
+  call void @llvm.spv.typedBufferStore(
+      target("spirv.Image", i32, 5, 2, 0, 0, 2, 0) %buffer0, i32 10, <4 x i32> %data)
+
+  ret void
+}
+
+attributes #0 = { convergent noinline norecurse "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+attributes #1 = { convergent noinline norecurse "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }