spirv: Add a postprocessing pass to fix up uses of OpSampledImage

SPIR-V requires that any instruction using the result of an
OpSampledImage instruction be in the same block as the OpSampledImage.
This is hard to guarantee in code generation but easy to fix after the
fact, by simply inserting a new OpSampledImage before the user of its
result if needed, with the new instruction having the same operands as
the original OpSampledImage.
This change adds a new pass to spv::Builder::postProcess that does this.
This might leave the original OpSampledImage instructions "orphaned"
with no users of their result ID, but dead code elimination would take
care of those further down the line.

SPIRV/SpvBuilder.h

@@ -855,6 +855,8 @@ public:
     void postProcess(Instruction&);
     // Hook to visit each non-32-bit sized float/int operation in a block.
     void postProcessType(const Instruction&, spv::Id typeId);
+    // move OpSampledImage instructions to be next to their users.
+    void postProcessSamplers();
 
     void dump(std::vector<unsigned int>&) const;
 

SPIRV/SpvPostProcess.cpp

@@ -483,6 +483,58 @@ void Builder::postProcessFeatures() {
     }
 }
 
+// SPIR-V requires that any instruction consuming the result of an OpSampledImage
+// be in the same block as the OpSampledImage instruction. This pass goes finds
+// uses of OpSampledImage where that is not the case and duplicates the
+// OpSampledImage to be immediately before the instruction that consumes it.
+// The old OpSampledImage is left in place, potentially with no users.
+void Builder::postProcessSamplers()
+{
+    // first, find all OpSampledImage instructions and store them in a map.
+    std::map<Id, Instruction*> sampledImageInstrs;
+    for (auto f: module.getFunctions()) {
+	for (auto b: f->getBlocks()) {
+	    for (auto &i: b->getInstructions()) {
+		if (i->getOpCode() == spv::OpSampledImage) {
+		    sampledImageInstrs[i->getResultId()] = i.get();
+		}
+	    }
+	}
+    }
+    // next find all uses of the given ids and rewrite them if needed.
+    for (auto f: module.getFunctions()) {
+	for (auto b: f->getBlocks()) {
+            auto &instrs = b->getInstructions();
+            for (size_t idx = 0; idx < instrs.size(); idx++) {
+                Instruction *i = instrs[idx].get();
+                for (int opnum = 0; opnum < i->getNumOperands(); opnum++) {
+                    // Is this operand of the current instruction the result of an OpSampledImage?
+                    if (i->isIdOperand(opnum) &&
+                        sampledImageInstrs.count(i->getIdOperand(opnum)))
+                    {
+                        Instruction *opSampImg = sampledImageInstrs[i->getIdOperand(opnum)];
+                        if (i->getBlock() != opSampImg->getBlock()) {
+                            Instruction *newInstr = new Instruction(getUniqueId(),
+                                                                    opSampImg->getTypeId(),
+                                                                    spv::OpSampledImage);
+                            newInstr->addIdOperand(opSampImg->getIdOperand(0));
+                            newInstr->addIdOperand(opSampImg->getIdOperand(1));
+                            newInstr->setBlock(b);
+
+                            // rewrite the user of the OpSampledImage to use the new instruction.
+                            i->setIdOperand(opnum, newInstr->getResultId());
+                            // insert the new OpSampledImage right before the current instruction.
+                            instrs.insert(instrs.begin() + idx,
+                                    std::unique_ptr<Instruction>(newInstr));
+                            idx++;
+                        }
+                    }
+                }
+            }
+	}
+    }
+}
+
 // comment in header
 void Builder::postProcess(bool compileOnly)
 {
@@ -491,6 +543,7 @@ void Builder::postProcess(bool compileOnly)
         postProcessCFG();
 
     postProcessFeatures();
+    postProcessSamplers();
 }
 
 }; // end spv namespace

SPIRV/spvIR.h

@@ -107,6 +107,14 @@ public:
         operands.push_back(id);
         idOperand.push_back(true);
     }
+    // This method is potentially dangerous as it can break assumptions
+    // about SSA and lack of forward references.
+    void setIdOperand(unsigned idx, Id id) {
+        assert(id);
+        assert(idOperand[idx]);
+        operands[idx] = id;
+    }
+
     void addImmediateOperand(unsigned int immediate) {
         operands.push_back(immediate);
         idOperand.push_back(false);
@@ -238,7 +246,7 @@ public:
     void addLocalVariable(std::unique_ptr<Instruction> inst) { localVariables.push_back(std::move(inst)); }
     const std::vector<Block*>& getPredecessors() const { return predecessors; }
     const std::vector<Block*>& getSuccessors() const { return successors; }
-    const std::vector<std::unique_ptr<Instruction> >& getInstructions() const {
+    std::vector<std::unique_ptr<Instruction> >& getInstructions() {
         return instructions;
     }
     const std::vector<std::unique_ptr<Instruction> >& getLocalVariables() const { return localVariables; }