spirv-opt: Fix when ConvertToHalfPass relaxes a image operand

Author: spencer-lunarg

source/opt/convert_to_half_pass.cpp

@@ -16,6 +16,8 @@
 
 #include "convert_to_half_pass.h"
 
+#include <cstdint>
+
 #include "source/opt/ir_builder.h"
 
 namespace spvtools {
@@ -195,6 +197,117 @@ bool ConvertToHalfPass::MatConvertCleanup(Instruction* inst) {
   return true;
 }
 
+// Return operand position of Image Operands <ID> or zero if there is none
+static constexpr uint32_t GetImageOperandsPosition(spv::Op opcode) {
+  uint32_t position = 0;
+  switch (opcode) {
+    case spv::Op::OpImageWrite:
+      return 3;
+    case spv::Op::OpImageSampleImplicitLod:
+    case spv::Op::OpImageSampleExplicitLod:
+    case spv::Op::OpImageSampleProjImplicitLod:
+    case spv::Op::OpImageSampleProjExplicitLod:
+    case spv::Op::OpImageFetch:
+    case spv::Op::OpImageRead:
+    case spv::Op::OpImageSparseSampleImplicitLod:
+    case spv::Op::OpImageSparseSampleExplicitLod:
+    case spv::Op::OpImageSparseSampleProjImplicitLod:
+    case spv::Op::OpImageSparseSampleProjExplicitLod:
+    case spv::Op::OpImageSparseFetch:
+    case spv::Op::OpImageSparseRead:
+      return 4;
+    case spv::Op::OpImageSampleDrefImplicitLod:
+    case spv::Op::OpImageSampleDrefExplicitLod:
+    case spv::Op::OpImageSampleProjDrefImplicitLod:
+    case spv::Op::OpImageSampleProjDrefExplicitLod:
+    case spv::Op::OpImageGather:
+    case spv::Op::OpImageDrefGather:
+    case spv::Op::OpImageSparseSampleDrefImplicitLod:
+    case spv::Op::OpImageSparseSampleDrefExplicitLod:
+    case spv::Op::OpImageSparseSampleProjDrefImplicitLod:
+    case spv::Op::OpImageSparseSampleProjDrefExplicitLod:
+    case spv::Op::OpImageSparseGather:
+    case spv::Op::OpImageSparseDrefGather:
+      return 5;
+    case spv::Op::OpImageSampleFootprintNV:
+      return 6;
+    default:
+      break;
+  }
+  return position;
+}
+
+bool ConvertToHalfPass::ImageOperandCleanup(Instruction* inst) {
+  if (image_ops_.count(inst->opcode()) == 0) return false;
+
+  const uint32_t image_operand_position =
+      GetImageOperandsPosition(inst->opcode());
+  // Image operands can be optional, also some things like |SignExtend| don't
+  // have additonal operands
+  if (image_operand_position == 0 ||
+      (image_operand_position + 1) >= inst->NumOperandWords()) {
+    return false;
+  }
+
+  bool modified = false;
+  const uint32_t mask = inst->GetSingleWordOperand(image_operand_position);
+  uint32_t next_operand = image_operand_position + 1;
+
+  InstructionBuilder builder(
+      context(), inst,
+      IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
+
+  uint32_t float32_type =
+      context()->get_type_mgr()->GetTypeInstruction(FloatScalarType(32));
+
+  // All mask after don't have a 32-bit restriction
+  for (uint32_t i = 1; i <= (uint32_t)spv::ImageOperandsMask::MinLod; i <<= 1) {
+    if ((mask & i) == 0) {
+      continue;
+    }
+
+    uint32_t update_operand = next_operand++;
+    const uint32_t ref_id = inst->GetSingleWordOperand(update_operand);
+    Instruction* ref_inst = get_def_use_mgr()->GetDef(ref_id);
+    Instruction* type_inst = get_def_use_mgr()->GetDef(ref_inst->type_id());
+    assert(type_inst->opcode() == spv::Op::OpTypeFloat);
+    if (type_inst->GetSingleWordInOperand(0) != 32) {
+      Instruction* cvt_inst =
+          builder.AddUnaryOp(float32_type, spv::Op::OpFConvert, ref_id);
+      if (cvt_inst == nullptr) {
+        status_ = Status::Failure;
+        return false;
+      }
+      inst->SetOperand(update_operand, {cvt_inst->result_id()});
+      get_def_use_mgr()->AnalyzeInstUse(inst);
+      modified = true;
+    }
+
+    // All masks have 1 operand, except grad has two
+    if (mask & uint32_t(spv::ImageOperandsMask::Grad)) {
+      uint32_t dy_update_operand = next_operand++;
+      const uint32_t dy_ref_id = inst->GetSingleWordOperand(dy_update_operand);
+      Instruction* dy_ref_inst = get_def_use_mgr()->GetDef(dy_ref_id);
+      Instruction* dy_type_inst =
+          get_def_use_mgr()->GetDef(dy_ref_inst->type_id());
+      assert(dy_type_inst->opcode() == spv::Op::OpTypeFloat);
+      if (dy_type_inst->GetSingleWordInOperand(0) != 32) {
+        Instruction* dy_cvt_inst =
+            builder.AddUnaryOp(float32_type, spv::Op::OpFConvert, ref_id);
+        if (dy_cvt_inst == nullptr) {
+          status_ = Status::Failure;
+          return false;
+        }
+        inst->SetOperand(dy_update_operand, {dy_cvt_inst->result_id()});
+        get_def_use_mgr()->AnalyzeInstUse(inst);
+        modified = true;
+      }
+    }
+  }
+
+  return modified;
+}
+
 bool ConvertToHalfPass::RemoveRelaxedDecoration(uint32_t id) {
   return context()->get_decoration_mgr()->RemoveDecorationsFrom(
       id, [](const Instruction& dec) {
@@ -452,6 +565,7 @@ bool ConvertToHalfPass::ProcessFunction(Function* func) {
         }
         for (auto ii = bb->begin(); ii != bb->end(); ++ii) {
           bool Mmodified = MatConvertCleanup(&*ii);
+          Mmodified |= ImageOperandCleanup(&*ii);
           if (status_ == Status::Failure) {
             ok = false;
             break;

source/opt/convert_to_half_pass.h

@@ -114,6 +114,10 @@ class ConvertToHalfPass : public Pass {
   // invalid so we need to clean them up.
   bool MatConvertCleanup(Instruction* inst);
 
+  // If |inst| has any image operands, make sure to covert it back to a 32-bit
+  // float
+  bool ImageOperandCleanup(Instruction* inst);
+
   // Call GenHalfInst on every instruction in |func|.
   // If code is generated for an instruction, replace the instruction
   // with the new instructions that are generated.