AdvancedInterfaceVariableScalarReplacement vectors handling fix

den-mentiei · den-mentiei · commit 3b397d053ad8 · 2025-08-06T13:57:12.000+02:00
It was crashing on this small shader:

#version 450
layout(location = 1) in vec3 x[2];
void main() { float crash = x[1].y; }

Because of non-supported OpAccessChain-s targeting the array-ed vectors' component.

Code now detects the component access and inserts a reduced OpAccessChain (targeting the replacement scalar variable component correctly). Added a test for this specific case.
Fixed copyright header I got wrong last time.
diff --git a/source/opt/adv_interface_var_sroa.cpp b/source/opt/adv_interface_var_sroa.cpp
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 Google LLC
+// Copyright (c) 2025 Epic Games, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -30,11 +30,20 @@ constexpr uint32_t kOpEntryPointInOperandInterface = 3;
 constexpr uint32_t kOpVariableStorageClassInOperandIndex = 0;
 constexpr uint32_t kOpTypeArrayElemTypeInOperandIndex = 0;
 constexpr uint32_t kOpTypeArrayLengthInOperandIndex = 1;
+constexpr uint32_t kOpTypeVectorComponentCountInOperandIndex = 1;
 constexpr uint32_t kOpTypeMatrixColCountInOperandIndex = 1;
 constexpr uint32_t kOpTypeMatrixColTypeInOperandIndex = 0;
 constexpr uint32_t kOpTypePtrTypeInOperandIndex = 1;
 constexpr uint32_t kOpConstantValueInOperandIndex = 0;
 
+// Get the component count of the OpTypeVector |vector_type|.
+uint32_t GetVectorComponentCount(Instruction* vector_type) {
+  assert(vector_type->opcode() == spv::Op::OpTypeVector);
+  uint32_t component_count =
+      vector_type->GetSingleWordInOperand(kOpTypeVectorComponentCountInOperandIndex);
+  return component_count;
+}
+
 // Get the length of the OpTypeArray |array_type|.
 uint32_t GetArrayLength(analysis::DefUseManager* def_use_mgr,
                         Instruction* array_type) {
@@ -223,6 +232,8 @@ Pass::Status AdvancedInterfaceVariableScalarReplacement::ProcessEntryPoint(
 
   ReplaceInEntryPoint(&entry_point, replaced_interface_vars, scalar_vars);
 
+  context()->InvalidateAnalysesExceptFor(IRContext::Analysis::kAnalysisNone);
+
   return status;
 }
 
@@ -330,19 +341,46 @@ bool AdvancedInterfaceVariableScalarReplacement::ReplaceInterfaceVariable(
     // We are going to replace the access chain with either direct usage of the
     // replacement scalar variable, or a set of composite loads/stores.
 
-    const Replacement* target =
+    LookupResult result =
         LookupReplacement(access_chain, &replacement, var.extra_array_length);
-    if (!target) {
+    if (!result.replacement) {
       // Error has been already logged by |LookupReplacement|.
       return false;
     }
+    const Replacement* target = result.replacement;
 
     if (!target->HasChildren() && var.extra_array_length == 0) {
-      // Replace with a direct use of the scalar variable.
       auto scalar = target->GetScalarVariable();
       assert(scalar);
-      context()->ReplaceAllUsesWith(access_chain->result_id(),
-                                    scalar->result_id());
+
+      uint32_t replacement = 0;
+      if (result.index >= 0) {
+        // Our scalar is a vector and access chain in question targets a
+        // specific component denoted by result.index.
+        assert(target->GetVectorComponentCount() > 0);
+        // Replace with an access chain into a direct use of the scalar variable.
+        uint32_t indirection_id = TakeNextId();
+        if (indirection_id == 0) {
+          return false;
+        }
+
+        uint32_t vector_component_type_id = context()->get_def_use_mgr()->GetDef(target->GetTypeId())->GetSingleWordInOperand(0);
+
+        uint32_t index_id = context()->get_constant_mgr()->GetUIntConstId(result.index);
+        Operand index_operand = {SPV_OPERAND_TYPE_ID, {index_id}};
+        std::unique_ptr<Instruction> vector_access_chain =
+            CreateAccessChain(context(), indirection_id, scalar,
+                              vector_component_type_id, index_operand);
+        replacement = vector_access_chain->result_id();
+
+        auto inst = access_chain->InsertBefore(std::move(vector_access_chain));
+        inst->UpdateDebugInfoFrom(access_chain);
+        get_def_use_mgr()->AnalyzeInstDef(inst);
+      } else {
+        // Replace with a direct use of the scalar variable.
+        replacement = scalar->result_id();
+      }
+      context()->ReplaceAllUsesWith(access_chain->result_id(), replacement);
     } else {
       // The current access chain's target is a composite, meaning that there
       // are other instructions using the pointer. We need to convert those to
@@ -732,7 +770,7 @@ bool AdvancedInterfaceVariableScalarReplacement::ReplaceStore(
   return true;
 }
 
-const AdvancedInterfaceVariableScalarReplacement::Replacement*
+AdvancedInterfaceVariableScalarReplacement::LookupResult
 AdvancedInterfaceVariableScalarReplacement::LookupReplacement(
     Instruction* access_chain, const Replacement* root,
     uint32_t extra_array_length) {
@@ -744,37 +782,59 @@ AdvancedInterfaceVariableScalarReplacement::LookupReplacement(
   // array, hence we skip it when looking-up the rest.
   uint32_t start_index = extra_array_length == 0 ? 1 : 2;
 
+  uint32_t num_indices = access_chain->NumInOperands();
+
   // Finds the target replacement, which might be a scalar or nested
   // composite.
-  for (uint32_t i = start_index; i < access_chain->NumInOperands(); ++i) {
+  for (uint32_t i = start_index; i < num_indices; ++i) {
     uint32_t index_id = access_chain->GetSingleWordInOperand(i);
 
     const analysis::Constant* index_constant =
         const_mgr->FindDeclaredConstant(index_id);
     if (!index_constant) {
       context()->EmitErrorMessage(
           "Variable cannot be replaced: index is not constant", access_chain);
-      return nullptr;
+      return {};
+    }
+
+    // OpAccessChain treats indices as signed.
+    int64_t index_value = index_constant->GetSignExtendedValue();
+
+    // Very last index can target the vector type, which we
+    // have as a scalar.
+    if (i == num_indices - 1) {
+      if (root->GetScalarVariable()) {
+        if (index_value < 0 ||
+            index_value >=
+                static_cast<int64_t>(root->GetVectorComponentCount())) {
+          // Out of bounds access, this is illegal IR.
+          // Notice that OpAccessChain indexing is 0-based, so we should also
+          // reject index == size-of-array.
+          context()->EmitErrorMessage(
+              "Variable cannot be replaced: invalid index", access_chain);
+          return {};
+        }
+        // Current root is our replacement scalar - a vector, in fact.
+        return {root, index_value};
+      }
     }
 
     assert(root->HasChildren());
     const auto& children = root->GetChildren();
 
-    // OpAccessChain treats indices as signed.
-    int64_t index_value = index_constant->GetSignExtendedValue();
     if (index_value < 0 ||
         index_value >= static_cast<int64_t>(children.size())) {
       // Out of bounds access, this is illegal IR.
       // Notice that OpAccessChain indexing is 0-based, so we should also
       // reject index == size-of-array.
       context()->EmitErrorMessage("Variable cannot be replaced: invalid index",
                                   access_chain);
-      return nullptr;
+      return {};
     }
 
     root = &children[index_value];
   }
-  return root;
+  return {root};
 }
 
 AdvancedInterfaceVariableScalarReplacement::Replacement
@@ -863,7 +923,12 @@ AdvancedInterfaceVariableScalarReplacement::CreateReplacementVariables(
       std::unique_ptr<Instruction> variable = CreateVariable(
           type->result_id(), storage_class, var.def, var.extra_array_length);
 
-      node->SetSingleScalarVariable(variable.get());
+      uint32_t vector_component_count = 0;
+      if (opcode == spv::Op::OpTypeVector) {
+        vector_component_count = GetVectorComponentCount(type);
+      }
+
+      node->SetSingleScalarVariable(variable.get(), vector_component_count);
       scalar_vars->push_back(variable.get());
 
       uint32_t var_id = variable->result_id();
diff --git a/source/opt/adv_interface_var_sroa.h b/source/opt/adv_interface_var_sroa.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 Google LLC
+// Copyright (c) 2025 Epic Games, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -90,14 +90,21 @@ class AdvancedInterfaceVariableScalarReplacement : public Pass {
 
     Instruction* GetScalarVariable() const { return scalar_var; }
 
-    void SetSingleScalarVariable(Instruction* var) { scalar_var = var; }
+    void SetSingleScalarVariable(Instruction* var, uint32_t in_vector_component_count) {
+      scalar_var = var;
+      vector_component_count = in_vector_component_count;
+    }
 
     uint32_t GetTypeId() const { return type_id; }
 
+    // Returns 0, if the Replacement is not a vector.
+    uint32_t GetVectorComponentCount() const { return vector_component_count; }
+
    private:
     std::vector<Replacement> children;
     Instruction* scalar_var;
     uint32_t type_id;
+    uint32_t vector_component_count;
   };
 
   // Collects all interface variables used by the |entry_point|.
@@ -171,14 +178,21 @@ class AdvancedInterfaceVariableScalarReplacement : public Pass {
                     Instruction* optional_access_chain,
                     uint32_t extra_array_length);
 
+  struct LookupResult {
+    // The replacement node, nullptr if not found.
+    const Replacement* replacement = nullptr;
+    // If |replacement| is a vector, which was also indexed by |access_chain|,
+    // this will have that used index value.
+    int64_t index = -1;
+  };
   // Looks up the replacement node according to the indices from the access
   // chain |access_chain|, using the passed |root| as a base. If any index in
   // the chain is non-constant or ouf-of-bound, return nullptr. If
   // |extra_array_length| is not zero, the first index in the chain is skipped,
   // as it is the one used for extra arrayness.
-  const Replacement* LookupReplacement(Instruction* access_chain,
-                                       const Replacement* root,
-                                       uint32_t extra_array_length);
+  LookupResult LookupReplacement(Instruction* access_chain,
+                                 const Replacement* root,
+                                 uint32_t extra_array_length);
 
   // Creates a variable with type |type_id| and storage class |storage_class|.
   // Debug info for the newly created variable is copied from the source
diff --git a/test/opt/adv_interface_var_sroa_test.cpp b/test/opt/adv_interface_var_sroa_test.cpp
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 Google LLC
+// Copyright (c) 2025 Epic Games, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -70,15 +70,11 @@ TEST_F(AdvancedInterfaceVariableScalarReplacementTest,
 ; CHECK-DAG: OpDecorate [[y]] Location 0
 ; CHECK-DAG: OpDecorate [[gl_InvocationID]] BuiltIn InvocationId
 ; CHECK-DAG: OpDecorate [[z0]] Location 0
-; CHECK-DAG: OpDecorate [[z0]] Component 0
 ; CHECK-DAG: OpDecorate [[z1]] Location 1
-; CHECK-DAG: OpDecorate [[z1]] Component 0
 ; CHECK-DAG: OpDecorate [[z0]] Patch
 ; CHECK-DAG: OpDecorate [[z1]] Patch
 ; CHECK-DAG: OpDecorate [[w0]] Location 2
-; CHECK-DAG: OpDecorate [[w0]] Component 0
 ; CHECK-DAG: OpDecorate [[w1]] Location 3
-; CHECK-DAG: OpDecorate [[w1]] Component 0
 ; CHECK-DAG: OpDecorate [[w0]] Patch
 ; CHECK-DAG: OpDecorate [[w1]] Patch
 ; CHECK-DAG: OpDecorate [[u0]] Location 3
@@ -335,6 +331,64 @@ TEST_F(AdvancedInterfaceVariableScalarReplacementTest,
   SinglePassRunAndMatch<AdvancedInterfaceVariableScalarReplacement>(spirv, true, true);
 }
 
+TEST_F(AdvancedInterfaceVariableScalarReplacementTest,
+       ReplaceInterfaceVarsWithScalas_Vectors) {
+  const std::string spirv = R"(
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Vertex %func "shader" %x
+
+; CHECK:     OpName [[y:%\w+]] "y"
+; CHECK:     OpName [[x0:%\w+]] "x[0]"
+; CHECK:     OpName [[x1:%\w+]] "x[1]"
+; CHECK-NOT: OpName {{%\w+}} "y"
+               OpName %x "x"
+               OpName %y "y"
+
+; CHECK-DAG: OpDecorate [[x0]] Location 1
+; CHECK-DAG: OpDecorate [[x1]] Location 2
+               OpDecorate %x Location 1
+
+      %float = OpTypeFloat 32
+        %int = OpTypeInt 32 1
+       %uint = OpTypeInt 32 0
+      %int_1 = OpConstant %int 1
+     %uint_1 = OpConstant %uint 1
+     %uint_2 = OpConstant %uint 2
+    %v3float = OpTypeVector %float 3
+%_arr_v3float_uint_2 = OpTypeArray %v3float %uint_2
+%_ptr_Input_float = OpTypePointer Input %float
+%_ptr_Input_v3float = OpTypePointer Input %v3float
+%_ptr_Input__arr_v3float_uint_2 = OpTypePointer Input %_arr_v3float_uint_2
+%_ptr_Function_float = OpTypePointer Function %float
+%_ptr_Function__vec3 = OpTypePointer Function %v3float
+
+          %x = OpVariable %_ptr_Input__arr_v3float_uint_2 Input
+; CHECK-DAG: [[x0]] = OpVariable %_ptr_Input_v3float Input
+; CHECK-DAG: [[x1]] = OpVariable %_ptr_Input_v3float Input
+
+     %void   = OpTypeVoid
+     %void_f = OpTypeFunction %void
+     %func   = OpFunction %void None %void_f
+     %label  = OpLabel
+
+          %y = OpVariable %_ptr_Function_float Function
+; CHECK-DAG [[y]] = OpVariable %_ptr_Function_float Function
+
+; CHECK: [[ptr:%\w+]] = OpAccessChain %_ptr_Input_float [[x1]] %uint_1
+; CHECK: [[val:%\w+]] = OpLoad %float [[ptr]]
+; CHECK:                OpStore [[y]] [[val]]
+    %x_z_ptr = OpAccessChain %_ptr_Input_float %x %int_1 %uint_1
+    %x_z_val = OpLoad %float %x_z_ptr
+               OpStore %y %x_z_val
+
+               OpReturn
+               OpFunctionEnd
+  )";
+
+  SinglePassRunAndMatch<AdvancedInterfaceVariableScalarReplacement>(spirv, true, true);
+}
+
 TEST_F(AdvancedInterfaceVariableScalarReplacementTest,
        CheckPatchDecorationPreservation) {
   // Make sure scalars for the variables with the extra arrayness have the extra
