[Fix, Test] (ShaderTranspiler): let the image-array select ladder carry an imageSize query, not just a read

MobileGL/MG_Test/ShaderTranspiler/LegalizeResourceArrayIndexTest.cpp

@@ -342,6 +342,21 @@ void main() {
     // An imageAtomic* reaches the array through OpImageTexelPointer, and running one per element
     // would perform every other element's atomic as well. The pass has to decline rather than
     // lower this.
+    // imageSize() on a dynamically indexed image array. The query carries the image in the same
+    // leading operand position as an imageLoad and answers with an int vector, so the select
+    // ladder spells it exactly - and unlike a read it touches no memory at all, so evaluating it
+    // for every element cannot even return undefined data.
+    constexpr const char* kUniformIndexedImageSizeQuery = R"(#version 450 core
+layout(local_size_x = 1) in;
+layout(rgba32f, binding = 0) uniform image2D g_image[4];
+layout(rgba32f, binding = 4) uniform image2DArray g_layered[2];
+layout(std430, binding = 8) buffer Out { ivec2 size; int layers; } g_out;
+uniform int g_index;
+void main() {
+    g_out.size = imageSize(g_image[g_index]);
+    g_out.layers = imageSize(g_layered[g_index]).z;
+}
+)";
     constexpr const char* kUniformIndexedImageAtomic = R"(#version 450 core
 layout(local_size_x = 1) in;
 layout(r32ui, binding = 0) uniform uimage2D g_image[4];
@@ -514,6 +529,36 @@ TEST(LegalizeResourceArrayIndexPass, LeavesADynamicallyIndexedSamplerArrayByteId
     EXPECT_EQ(output, input);
 }
 
+// imageSize() on a dynamically indexed image array used to lose the whole stage: the consumer
+// whitelist accepted only OpImageRead/OpImageWrite, so the chain was declined and the illegal
+// subscript reached the ES compiler intact. It is the same select ladder as a read - the query
+// takes the image in in-operand 0 and produces an int vector - and it reads no memory, so the
+// elements the shader did not ask for cost nothing but the instruction.
+TEST(LegalizeResourceArrayIndexPass, LowersAUniformIndexedImageSizeQueryToSelects) {
+    const Vector<Uint32> input = CompileCompute(kUniformIndexedImageSizeQuery);
+    ASSERT_FALSE(input.empty());
+    EXPECT_TRUE(HasDynamicImageArrayIndex(input));
+    EXPECT_EQ(CountOpcode(input, spv::Op::OpImageQuerySize), 2u);
+    EXPECT_EQ(CountOpcode(input, spv::Op::OpSelect), 0u);
+
+    Vector<Uint32> output;
+    ASSERT_TRUE(ShaderCompiler::LegalizeResourceArrayIndexingForEssl(input, output, true));
+    ASSERT_FALSE(output.empty());
+    EXPECT_FALSE(HasDynamicImageArrayIndex(output));
+    // Four elements for g_image and two for g_layered, one query apiece, and one select per
+    // element past the first of each ladder.
+    EXPECT_EQ(CountOpcode(output, spv::Op::OpImageQuerySize), 6u);
+    EXPECT_EQ(CountOpcode(output, spv::Op::OpSelect), 4u);
+    EXPECT_EQ(CountOpcode(output, spv::Op::OpSwitch), 0u) << "a query produces a value, so no control flow";
+    EXPECT_TRUE(Validates(output));
+
+    const EsslAttempt after = EmitEssl(output);
+    ASSERT_TRUE(after.succeeded) << after.error;
+    for (int element = 0; element < 4; ++element) {
+        EXPECT_NE(after.text.find("g_image[" + std::to_string(element) + "]"), String::npos) << after.text;
+    }
+}
+
 // An imageAtomic* is the shape the lowering must refuse: its per-element rebuild would run every
 // other element's read-modify-write. Declining leaves the illegal subscript in place - which is
 // what the latched warning in LegalizeResourceArrayIndexingForEssl is for - but a half-transform

MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.cpp

@@ -737,12 +737,18 @@ namespace MobileGL {
                         return;
                     case spv::Op::OpImageWrite:
                     case spv::Op::OpImageRead:
+                    // imageSize()/imageSamples() carry the image in the same leading operand
+                    // position as an OpImageRead and produce an int or int vector, so the same
+                    // select ladder rebuilds them exactly - and a size query touches no memory
+                    // at all, which makes evaluating it for every element strictly safer than
+                    // the read the ladder was written for.
+                    case spv::Op::OpImageQuerySize:
+                    case spv::Op::OpImageQuerySizeLod:
                         if (consumer == nullptr) consumer = user;
                         return;
                     default:
-                        // A sampled-image construction, a query, a copy, an argument to a
+                        // A sampled-image construction, a copy, an argument to a function:
-                        // function: shapes whose per-element rebuild this pass cannot spell
+                        // shapes whose per-element rebuild this pass cannot spell exactly.
-                        // exactly.
                         unsupportedConsumer = true;
                         return;
                     }
@@ -762,7 +768,7 @@ namespace MobileGL {
 
                 return consumer->opcode() == spv::Op::OpImageWrite
                            ? LowerImageWrite(accessChain, arrayLength, load, consumer)
-                           : LowerImageRead(accessChain, arrayLength, load, consumer);
+                           : LowerImageReadOrQuery(accessChain, arrayLength, load, consumer);
             }
 
             // Drops |load| and |accessChain| once the rewrite above has taken their last user,
@@ -865,38 +871,42 @@ namespace MobileGL {
                 return LoweringOutcome::Changed;
             }
 
-            // A read needs no control flow: read every element through a constant index and pick
+            // A value-producing consumer - an OpImageRead, or an imageSize()/imageSamples() query -
+            // needs no control flow: run it against every element through a constant index and pick
             // with OpSelect. The selection happens on the RESULT, not on the image object - an
             // opaque type may not be selected at all (pre-1.4 OpSelect takes pointers, scalars
             // and vectors only, and ESSL has no ternary on an image), so what is duplicated is
-            // the OpImageRead.
+            // the consuming instruction itself. Every such consumer carries the image in in-operand
+            // 0 and nothing else that is per-element, so one rebuild spells all of them.
             //
-            // Reading the elements the shader did not ask for is safe: every one of them is an
+            // Running the elements the shader did not ask for is safe: every one of them is an
             // image this stage already declares, and GL 4.6 7.11.2 makes a load through an
             // image unit whose binding is missing or incompatible return undefined DATA - never
-            // an error, and never a fault - which the select then discards. Contrast an
+            // an error, and never a fault - which the select then discards. A size query does not
-            // imageAtomic*, which LowerImageChain refuses for exactly the opposite reason.
+            // even touch memory. Contrast an imageAtomic*, which LowerImageChain refuses for
+            // exactly the opposite reason.
             LegalizeResourceArrayIndexPass::LoweringOutcome
-            LegalizeResourceArrayIndexPass::LowerImageRead(Instruction* accessChain, uint32_t arrayLength,
+            LegalizeResourceArrayIndexPass::LowerImageReadOrQuery(Instruction* accessChain,
-                                                           Instruction* load, Instruction* imageRead) {
+                                                                  uint32_t arrayLength, Instruction* load,
+                                                                  Instruction* consumer) {
                 auto* irContext = context();
                 uint32_t conditionTypeId = 0;
                 uint32_t dimension = 0;
-                if (!TryGetSelectConditionType(irContext, imageRead->type_id(), &conditionTypeId,
+                if (!TryGetSelectConditionType(irContext, consumer->type_id(), &conditionTypeId,
                                                &dimension)) {
-                    MGLOG_D("[spirv] image array index: read type is not selectable, declining");
+                    MGLOG_D("[spirv] image array index: result type is not selectable, declining");
                     return LoweringOutcome::Declined;
                 }
                 const uint32_t boolTypeId = irContext->get_type_mgr()->GetBoolTypeId();
                 const uint32_t indexId = accessChain->GetSingleWordInOperand(1);
                 const uint32_t imageTypeId = load->type_id();
                 std::vector<Operand> tailOperands;
-                for (uint32_t i = 1; i < imageRead->NumInOperands(); ++i) {
+                for (uint32_t i = 1; i < consumer->NumInOperands(); ++i) {
-                    tailOperands.push_back(imageRead->GetInOperand(i));
+                    tailOperands.push_back(consumer->GetInOperand(i));
                 }
 
                 InstructionBuilder builder(
-                    irContext, imageRead,
+                    irContext, consumer,
                     IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
 
                 uint32_t selectedId = 0;
@@ -906,18 +916,18 @@ namespace MobileGL {
                         CloneChainWithConstantIndex(builder, irContext, accessChain, constantId);
                     Instruction* elementImage = builder.AddLoad(imageTypeId, elementChain->result_id());
 
-                    std::vector<Operand> readOperands;
+                    std::vector<Operand> elementOperands;
-                    readOperands.push_back({SPV_OPERAND_TYPE_ID, {elementImage->result_id()}});
+                    elementOperands.push_back({SPV_OPERAND_TYPE_ID, {elementImage->result_id()}});
                     for (const Operand& tailOperand : tailOperands) {
-                        readOperands.push_back(tailOperand);
+                        elementOperands.push_back(tailOperand);
                     }
-                    Instruction* elementRead = builder.AddInstruction(
+                    Instruction* elementResult = builder.AddInstruction(
-                        MakeUnique<Instruction>(irContext, spv::Op::OpImageRead, imageRead->type_id(),
+                        MakeUnique<Instruction>(irContext, consumer->opcode(), consumer->type_id(),
-                                                irContext->TakeNextId(), readOperands));
+                                                irContext->TakeNextId(), elementOperands));
                     if (element == 0) {
                         // Element 0 is the else-arm of the whole ladder, so an out-of-range
                         // index reads it - an undefined element for an undefined index.
-                        selectedId = elementRead->result_id();
+                        selectedId = elementResult->result_id();
                         continue;
                     }
 
@@ -929,17 +939,17 @@ namespace MobileGL {
                         conditionId = builder.AddCompositeConstruct(conditionTypeId, components)->result_id();
                     }
                     selectedId = builder
-                                     .AddSelect(imageRead->type_id(), conditionId,
+                                     .AddSelect(consumer->type_id(), conditionId,
-                                                elementRead->result_id(), selectedId)
+                                                elementResult->result_id(), selectedId)
                                      ->result_id();
                 }
 
-                irContext->ReplaceAllUsesWith(imageRead->result_id(), selectedId);
+                irContext->ReplaceAllUsesWith(consumer->result_id(), selectedId);
-                irContext->KillInst(imageRead);
+                irContext->KillInst(consumer);
                 KillImageChainIfDead(accessChain, load);
                 irContext->InvalidateAnalysesExceptFor(IRContext::kAnalysisNone);
-                MGLOG_D("[spirv] image array index: lowered a dynamic imageLoad to %u constant-indexed "
+                MGLOG_D("[spirv] image array index: lowered a dynamic image read/query to %u "
-                        "reads",
+                        "constant-indexed operations",
                         arrayLength);
                 return LoweringOutcome::Changed;
             }

MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.h

@@ -146,9 +146,10 @@ namespace MobileGL {
                 LoweringOutcome LowerImageWrite(spvtools::opt::Instruction* accessChain, uint32_t arrayLength,
                                                 spvtools::opt::Instruction* load,
                                                 spvtools::opt::Instruction* imageWrite);
-                LoweringOutcome LowerImageRead(spvtools::opt::Instruction* accessChain, uint32_t arrayLength,
+                LoweringOutcome LowerImageReadOrQuery(spvtools::opt::Instruction* accessChain,
-                                               spvtools::opt::Instruction* load,
+                                                      uint32_t arrayLength,
-                                               spvtools::opt::Instruction* imageRead);
+                                                      spvtools::opt::Instruction* load,
+                                                      spvtools::opt::Instruction* consumer);
 
                 Mode m_mode;
             };