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 @@ class Builder {
     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 @@ class Instruction {
         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 @@ class Block {
     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; }

Test/baseResults/spv.sampledImageBlock.frag.out

@@ -0,0 +1,94 @@
+spv.sampledImageBlock.frag
+// Module Version 10000
+// Generated by (magic number): 8000b
+// Id's are bound by 55
+
+                              Capability Shader
+               1:             ExtInstImport  "GLSL.std.450"
+                              MemoryModel Logical GLSL450
+                              EntryPoint Fragment 4  "main" 36 45
+                              ExecutionMode 4 OriginUpperLeft
+                              Source GLSL 450
+                              Name 4  "main"
+                              Name 9  "texel"
+                              Name 12  "tex0"
+                              Name 16  "samp0"
+                              Name 21  "ParamBuffer"
+                              MemberName 21(ParamBuffer) 0  "cond"
+                              Name 23  "paramBuffer"
+                              Name 36  "texCoord"
+                              Name 45  "fragColor"
+                              Decorate 12(tex0) DescriptorSet 0
+                              Decorate 12(tex0) Binding 0
+                              Decorate 16(samp0) DescriptorSet 0
+                              Decorate 16(samp0) Binding 1
+                              MemberDecorate 21(ParamBuffer) 0 Offset 0
+                              Decorate 21(ParamBuffer) Block
+                              Decorate 23(paramBuffer) DescriptorSet 0
+                              Decorate 23(paramBuffer) Binding 2
+                              Decorate 36(texCoord) Flat
+                              Decorate 36(texCoord) Location 0
+                              Decorate 45(fragColor) Location 0
+               2:             TypeVoid
+               3:             TypeFunction 2
+               6:             TypeFloat 32
+               7:             TypeVector 6(float) 4
+               8:             TypePointer Function 7(fvec4)
+              10:             TypeImage 6(float) 2D sampled format:Unknown
+              11:             TypePointer UniformConstant 10
+        12(tex0):     11(ptr) Variable UniformConstant
+              14:             TypeSampler
+              15:             TypePointer UniformConstant 14
+       16(samp0):     15(ptr) Variable UniformConstant
+              18:             TypeSampledImage 10
+              20:             TypeInt 32 1
+ 21(ParamBuffer):             TypeStruct 20(int)
+              22:             TypePointer Uniform 21(ParamBuffer)
+ 23(paramBuffer):     22(ptr) Variable Uniform
+              24:     20(int) Constant 0
+              25:             TypePointer Uniform 20(int)
+              28:             TypeBool
+              30:             TypeVector 20(int) 2
+              31:             TypePointer Function 30(ivec2)
+              35:             TypePointer Input 30(ivec2)
+    36(texCoord):     35(ptr) Variable Input
+              44:             TypePointer Output 7(fvec4)
+   45(fragColor):     44(ptr) Variable Output
+              46:             TypeVector 6(float) 3
+              49:    6(float) Constant 1065353216
+         4(main):           2 Function None 3
+               5:             Label
+        9(texel):      8(ptr) Variable Function
+              32:     31(ptr) Variable Function
+              13:          10 Load 12(tex0)
+              17:          14 Load 16(samp0)
+              19:          18 SampledImage 13 17
+              26:     25(ptr) AccessChain 23(paramBuffer) 24
+              27:     20(int) Load 26
+              29:    28(bool) IEqual 27 24
+                              SelectionMerge 34 None
+                              BranchConditional 29 33 38
+              33:               Label
+              37:   30(ivec2)   Load 36(texCoord)
+                                Store 32 37
+                                Branch 34
+              38:               Label
+              39:   30(ivec2)   Load 36(texCoord)
+              40:   30(ivec2)   VectorShuffle 39 39 1 0
+                                Store 32 40
+                                Branch 34
+              34:             Label
+              41:   30(ivec2) Load 32
+              54:          18 SampledImage 13 17
+              42:          10 Image 54
+              43:    7(fvec4) ImageFetch 42 41 Lod 24
+                              Store 9(texel) 43
+              47:    7(fvec4) Load 9(texel)
+              48:   46(fvec3) VectorShuffle 47 47 0 1 2
+              50:    6(float) CompositeExtract 48 0
+              51:    6(float) CompositeExtract 48 1
+              52:    6(float) CompositeExtract 48 2
+              53:    7(fvec4) CompositeConstruct 50 51 52 49
+                              Store 45(fragColor) 53
+                              Return
+                              FunctionEnd

Test/spv.sampledImageBlock.frag

@@ -0,0 +1,21 @@
+#version 450
+
+layout(set = 0, binding = 0) uniform texture2D tex0;
+layout(set = 0, binding = 1) uniform sampler samp0;
+layout(set = 0, binding = 2) uniform ParamBuffer {
+    int cond;
+} paramBuffer;
+
+layout(location = 0) out vec4 fragColor;
+layout(location = 0) in flat ivec2 texCoord;
+
+void main() {
+    // get input
+
+    const vec4 texel = texelFetch(sampler2D(tex0, samp0),
+            paramBuffer.cond == 0 ? texCoord.xy : texCoord.yx,
+            0);
+
+    fragColor = vec4(texel.xyz, 1.0);
+
+}

gtests/Spv.FromFile.cpp

@@ -551,6 +551,7 @@ INSTANTIATE_TEST_SUITE_P(
         "spv.ext.textureShadowLod.error.frag",
         "spv.floatFetch.frag",
         "spv.atomicRvalue.error.vert",
+        "spv.sampledImageBlock.frag",
     })),
     FileNameAsCustomTestSuffix
 );