[dxil2spv] Translate extractValue and bufferStore (#4404)

Add support for translation of LLVM extractValue operations and DXIL
bufferStore operations to SPIR-V instructions.

The bufferStore implementation is limited to what is required to
support a simple passthrough compute shader. For the yet unsupported
cases, an error is emitted.

tools/clang/test/Dxil2Spv/passthru-cs.ll

@@ -86,7 +86,7 @@ attributes #2 = { nounwind }
 ; ; SPIR-V
 ; ; Version: 1.0
 ; ; Generator: Google spiregg; 0
-; ; Bound: 56
+; ; Bound: 59
 ; ; Schema: 0
 ;                OpCapability Shader
 ;                OpMemoryModel Logical GLSL450
@@ -167,8 +167,9 @@ attributes #2 = { nounwind }
 ;          %54 = OpAccessChain %_ptr_Function_int %24 %uint_4
 ;          %55 = OpBitcast %int %53
 ;                OpStore %54 %55
+;          %56 = OpAccessChain %_ptr_Function_int %24 %uint_0
+;          %57 = OpAccessChain %_ptr_Uniform_uint %4 %uint_0 %28
+;          %58 = OpBitcast %uint %56
+;                OpStore %57 %58
 ;                OpReturn
 ;                OpFunctionEnd
-; CHECK-ERRORS:
-; error: Unhandled LLVM instruction:   %6 = extractvalue %dx.types.ResRet.i32 %5, 0
-; error: Unhandled DXIL opcode: BufferStore

tools/clang/tools/dxil2spv/lib/dxil2spv.cpp

@@ -10,6 +10,7 @@
 
 #include "dxil2spv.h"
 
+#include "dxc/DXIL/DxilConstants.h"
 #include "dxc/DXIL/DxilModule.h"
 #include "dxc/DXIL/DxilUtil.h"
 #include "dxc/DxilContainer/DxilContainer.h"
@@ -19,6 +20,8 @@
 
 #include "clang/SPIRV/SpirvInstruction.h"
 #include "clang/SPIRV/SpirvType.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
@@ -274,6 +277,9 @@ void Translator::createInstruction(llvm::Instruction &instruction) {
     case hlsl::DXIL::OpCode::BufferLoad: {
       createBufferLoadInstruction(callInstruction);
     } break;
+    case hlsl::DXIL::OpCode::BufferStore: {
+      createBufferStoreInstruction(callInstruction);
+    } break;
     default: {
       emitError("Unhandled DXIL opcode: %0")
           << hlsl::OP::GetOpCodeName(dxilOpcode);
@@ -286,6 +292,12 @@ void Translator::createInstruction(llvm::Instruction &instruction) {
         cast<llvm::BinaryOperator>(instruction);
     createBinaryOpInstruction(binaryInstruction);
   }
+  // Handle extractvalue instructions.
+  else if (isa<llvm::ExtractValueInst>(instruction)) {
+    llvm::ExtractValueInst &extractValueInstruction =
+        cast<llvm::ExtractValueInst>(instruction);
+    createExtractValueInstruction(extractValueInstruction);
+  }
   // Handle return instructions.
   else if (isa<llvm::ReturnInst>(instruction)) {
     spvBuilder.createReturn({});
@@ -485,6 +497,14 @@ void Translator::createBufferLoadInstruction(llvm::CallInst &instruction) {
       getSpirvInstruction(instruction.getArgOperand(
           hlsl::DXIL::OperandIndex::kBufferLoadCoord0OpIdx));
 
+  // TODO: Handle coordinate c1 if defined.
+  if (!isa<llvm::UndefValue>(instruction.getArgOperand(
+          hlsl::DXIL::OperandIndex::kBufferLoadCoord1OpIdx))) {
+    emitError("BufferLoad not yet implemented to handle given coordinates: %0",
+              instruction);
+    return;
+  }
+
   // Initialize each field in the struct.
   for (size_t i = 0; i < structType->getFields().size(); i++) {
     // Add offset for current field.
@@ -515,6 +535,94 @@ void Translator::createBufferLoadInstruction(llvm::CallInst &instruction) {
   instructionMap[&instruction] = returnVar;
 }
 
+void Translator::createBufferStoreInstruction(llvm::CallInst &instruction) {
+  // TODO: Extend this function to work with all buffer types on which it is
+  // used, not just ByteAddressBuffers.
+
+  // ByteAddressBuffers are represented as a struct with one member that is a
+  // runtime array of unsigned integers. The SPIR-V OpAccessChain instruction is
+  // then used to access that offset, and OpStore is used to store integers
+  // into the array.
+
+  // clang-format off
+  // For example, the following DXIL instruction:
+  //   call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle %res, i32 %c0, i32 undef, i32 %v0, i32 undef, i32 undef, i32 undef, i8 1)
+  // would translate to the following SPIR-V instructions:
+  //   %x = OpAccessChain %_ptr_Uniform_uint %res %uint_0 %c0
+  //   %y = OpBitcast %uint %v0
+  //         OpStore %x %y
+  // clang-format on
+
+  // Get module output variable corresponding to given DXIL handle.
+  spirv::SpirvInstruction *outputVar =
+      getSpirvInstruction(instruction.getArgOperand(
+          hlsl::DXIL::OperandIndex::kBufferStoreHandleOpIdx));
+
+  // Translate indices into resource buffer to SPIR-V instructions.
+  auto uint32 = spvContext.getUIntType(32);
+  spirv::SpirvConstant *indexIntoStruct =
+      spvBuilder.getConstantInt(uint32, llvm::APInt(32, 0));
+  spirv::SpirvInstruction *index =
+      getSpirvInstruction(instruction.getArgOperand(
+          hlsl::DXIL::OperandIndex::kBufferStoreCoord0OpIdx));
+
+  // TODO: Handle coordinate c1 if defined.
+  if (!isa<llvm::UndefValue>(instruction.getArgOperand(
+          hlsl::DXIL::OperandIndex::kBufferStoreCoord1OpIdx))) {
+    emitError("BufferStore not yet implemented to handle given coordinates: %0",
+              instruction);
+    return;
+  }
+
+  // TODO: Handle other masks and values.
+  if (cast<llvm::ConstantInt>(
+          instruction.getArgOperand(
+              hlsl::DXIL::OperandIndex::kBufferStoreMaskOpIdx))
+          ->getLimitedValue() != 1) {
+    emitError(
+        "BufferStore not yet implemented to handle given mask and values: %0",
+        instruction);
+    return;
+  }
+
+  // Create access chain and store.
+  spirv::SpirvAccessChain *outputVarPtr = spvBuilder.createAccessChain(
+      uint32, outputVar, {indexIntoStruct, index}, {});
+  spirv::SpirvInstruction *valueToStore =
+      getSpirvInstruction(instruction.getArgOperand(
+          hlsl::DXIL::OperandIndex::kBufferStoreVal0OpIdx));
+  // LLVM types are signless, so type conversions are not 1-to-1. A bitcast on
+  // the value may be necessary before storing.
+  if (valueToStore->getResultType() != uint32)
+    valueToStore =
+        spvBuilder.createUnaryOp(spv::Op::OpBitcast, uint32, valueToStore, {});
+  spvBuilder.createStore(outputVarPtr, valueToStore, {});
+}
+
+void Translator::createExtractValueInstruction(
+    llvm::ExtractValueInst &instruction) {
+  // Get corresponding SPIR-V intstruction for given aggregate value.
+  spirv::SpirvInstruction *spvInstruction =
+      getSpirvInstruction(instruction.getAggregateOperand());
+
+  // Translate DXIL indices into aggregate to SPIR-V instructions.
+  auto uint32 = spvContext.getUIntType(32);
+  std::vector<spirv::SpirvInstruction *> indices;
+  indices.reserve(instruction.getNumIndices());
+  for (unsigned index : instruction.indices()) {
+    spirv::SpirvConstant *spvIndex =
+        spvBuilder.getConstantInt(uint32, llvm::APInt(32, index));
+    indices.push_back(spvIndex);
+  }
+
+  // Create access chain and save mapping.
+  const spirv::SpirvType *returnType = toSpirvType(instruction.getType());
+  spirv::SpirvAccessChain *accessChain =
+      spvBuilder.createAccessChain(returnType, spvInstruction, indices, {});
+
+  instructionMap[&instruction] = accessChain;
+}
+
 bool Translator::spirvToolsValidate(std::vector<uint32_t> *mod,
                                     std::string *messages) {
   spvtools::SpirvTools tools(featureManager.getTargetEnv());

tools/clang/tools/dxil2spv/lib/dxil2spv.h

@@ -83,6 +83,8 @@ private:
   void createBinaryOpInstruction(llvm::BinaryOperator &instruction);
   void createHandleInstruction(llvm::CallInst &instruction);
   void createBufferLoadInstruction(llvm::CallInst &instruction);
+  void createBufferStoreInstruction(llvm::CallInst &instruction);
+  void createExtractValueInstruction(llvm::ExtractValueInst &instruction);
 
   // SPIR-V Tools wrapper functions.
   bool spirvToolsValidate(std::vector<uint32_t> *mod, std::string *messages);