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

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
 );