diff --git a/src/compiler/nir/nir_linking_helpers.c b/src/compiler/nir/nir_linking_helpers.c
index 8d7302c93c1..e1cf58f3514 100644
--- a/src/compiler/nir/nir_linking_helpers.c
+++ b/src/compiler/nir/nir_linking_helpers.c
@@ -1007,9 +1007,6 @@ nir_assign_io_var_locations(struct exec_list *var_list, unsigned *size,
 
    sort_varyings(var_list);
 
-   const int base = stage == MESA_SHADER_FRAGMENT ?
-      (int) FRAG_RESULT_DATA0 : (int) VARYING_SLOT_VAR0;
-
    int UNUSED last_loc = 0;
    bool last_partial = false;
    nir_foreach_variable(var, var_list) {
@@ -1019,6 +1016,15 @@ nir_assign_io_var_locations(struct exec_list *var_list, unsigned *size,
          type = glsl_get_array_element(type);
       }
 
+      int base;
+      if (var->data.mode == nir_var_shader_in && stage == MESA_SHADER_VERTEX)
+         base = VERT_ATTRIB_GENERIC0;
+      else if (var->data.mode == nir_var_shader_out &&
+               stage == MESA_SHADER_FRAGMENT)
+         base = FRAG_RESULT_DATA0;
+      else
+         base = VARYING_SLOT_VAR0;
+
       unsigned var_size;
       if (var->data.compact) {
          /* compact variables must be arrays of scalars */