nir: Make nir_ssa_def_rewrite_uses_after take an SSA value

This replaces the new_src parameter of nir_ssa_def_rewrite_uses_after()
with an SSA def, and rewrites all the users as needed.

Acked-by: Alyssa Rosenzweig <alyssa@collabora.com>
Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Reviewed-by: Eric Anholt <eric@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/9383>

src/broadcom/compiler/v3d_nir_lower_image_load_store.c

@@ -186,7 +186,7 @@ v3d_nir_lower_image_load(nir_builder *b, nir_intrinsic_instr *instr)
                               nir_unpack_half_2x16_split_y(b, ba));
         }
 
-        nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, nir_src_for_ssa(result),
+        nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, result,
                                        result->parent_instr);
 }
 

src/compiler/glsl/gl_nir_lower_buffers.c

@@ -247,7 +247,7 @@ lower_buffer_interface_derefs_impl(nir_function_impl *impl,
                   intrin->dest.ssa.bit_size = 32;
                   nir_ssa_def *bval = nir_i2b(&b, &intrin->dest.ssa);
                   nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                                 nir_src_for_ssa(bval),
+                                                 bval,
                                                  bval->parent_instr);
                   progress = true;
                }

src/compiler/nir/nir.c

@@ -1407,10 +1407,10 @@ is_instr_between(nir_instr *start, nir_instr *end, nir_instr *between)
  * def->parent_instr and that after_me comes after def->parent_instr.
  */
 void
-nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_src new_src,
+nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_ssa_def *new_ssa,
                                nir_instr *after_me)
 {
-   if (new_src.is_ssa && def == new_src.ssa)
+   if (def == new_ssa)
       return;
 
    nir_foreach_use_safe(use_src, def) {
@@ -1420,11 +1420,14 @@ nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_src new_src,
        * the instruction list.
        */
       if (!is_instr_between(def->parent_instr, after_me, use_src->parent_instr))
-         nir_instr_rewrite_src(use_src->parent_instr, use_src, new_src);
+         nir_instr_rewrite_src_ssa(use_src->parent_instr, use_src, new_ssa);
    }
 
-   nir_foreach_if_use_safe(use_src, def)
+   nir_foreach_if_use_safe(use_src, def) {
-      nir_if_rewrite_condition(use_src->parent_if, new_src);
+      nir_if_rewrite_condition_ssa(use_src->parent_if,
+                                   &use_src->parent_if->condition,
+                                   new_ssa);
+   }
 }
 
 nir_component_mask_t

src/compiler/nir/nir.h

@@ -3947,7 +3947,7 @@ nir_ssa_dest_init_for_type(nir_instr *instr, nir_dest *dest,
 }
 void nir_ssa_def_rewrite_uses(nir_ssa_def *def, nir_ssa_def *new_ssa);
 void nir_ssa_def_rewrite_uses_src(nir_ssa_def *def, nir_src new_src);
-void nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_src new_src,
+void nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_ssa_def *new_ssa,
                                     nir_instr *after_me);
 
 nir_component_mask_t nir_ssa_def_components_read(const nir_ssa_def *def);

src/compiler/nir/nir_deref.c

@@ -1251,7 +1251,7 @@ opt_load_vec_deref(nir_builder *b, nir_intrinsic_instr *load)
          data = nir_bitcast_vector(b, &load->dest.ssa, old_bit_size);
       data = resize_vector(b, data, old_num_comps);
 
-      nir_ssa_def_rewrite_uses_after(&load->dest.ssa, nir_src_for_ssa(data),
+      nir_ssa_def_rewrite_uses_after(&load->dest.ssa, data,
                                      data->parent_instr);
       return true;
    }

src/compiler/nir/nir_lower_array_deref_of_vec.c

@@ -155,7 +155,7 @@ nir_lower_array_deref_of_vec_impl(nir_function_impl *impl,
                nir_instr_remove(&intrin->instr);
             } else {
                nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                              nir_src_for_ssa(scalar),
+                                              scalar,
                                               scalar->parent_instr);
             }
             progress = true;

src/compiler/nir/nir_lower_bit_size.c

@@ -205,7 +205,7 @@ lower_phi_instr(nir_builder *b, nir_phi_instr *phi, unsigned bit_size,
    b->cursor = nir_after_instr(&last_phi->instr);
 
    nir_ssa_def *new_dest = nir_u2u(b, &phi->dest.ssa, old_bit_size);
-   nir_ssa_def_rewrite_uses_after(&phi->dest.ssa, nir_src_for_ssa(new_dest),
+   nir_ssa_def_rewrite_uses_after(&phi->dest.ssa, new_dest,
                                   new_dest->parent_instr);
 }
 

src/compiler/nir/nir_lower_io_to_vector.c

@@ -481,7 +481,7 @@ nir_lower_io_to_vector_impl(nir_function_impl *impl, nir_variable_mode modes)
             nir_ssa_def *new_vec = nir_channels(&b, &intrin->dest.ssa,
                                                 vec4_comp_mask >> new_frac);
             nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                           nir_src_for_ssa(new_vec),
+                                           new_vec,
                                            new_vec->parent_instr);
 
             progress = true;
@@ -622,7 +622,7 @@ nir_vectorize_tess_levels_impl(nir_function_impl *impl)
             b.cursor = nir_after_instr(instr);
             nir_ssa_def *val = &intrin->dest.ssa;
             nir_ssa_def *comp = nir_channel(&b, val, index);
-            nir_ssa_def_rewrite_uses_after(val, nir_src_for_ssa(comp), comp->parent_instr);
+            nir_ssa_def_rewrite_uses_after(val, comp, comp->parent_instr);
          }
 
          progress = true;

src/compiler/nir/nir_lower_pntc_ytransform.c

@@ -79,7 +79,7 @@ lower_load_pointcoord(lower_pntc_ytransform_state *state,
                                         nir_channel(b, pntc, 0),
                                         nir_fadd(b, offset, scaled));
 
-   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(flipped_pntc),
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, flipped_pntc,
                                   flipped_pntc->parent_instr);
 }
 

src/compiler/nir/nir_lower_tex.c

@@ -843,7 +843,7 @@ swizzle_tg4_broadcom(nir_builder *b, nir_tex_instr *tex)
    unsigned swiz[4] = { 2, 3, 1, 0 };
    nir_ssa_def *swizzled = nir_swizzle(b, &tex->dest.ssa, swiz, 4);
 
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(swizzled),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, swizzled,
                                   swizzled->parent_instr);
 }
 
@@ -883,7 +883,7 @@ swizzle_result(nir_builder *b, nir_tex_instr *tex, const uint8_t swizzle[4])
       }
    }
 
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(swizzled),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, swizzled,
                                   swizzled->parent_instr);
 }
 
@@ -906,7 +906,7 @@ linearize_srgb_result(nir_builder *b, nir_tex_instr *tex)
                                   nir_channel(b, rgb, 2),
                                   nir_channel(b, &tex->dest.ssa, 3));
 
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(result),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, result,
                                   result->parent_instr);
 }
 
@@ -983,7 +983,7 @@ lower_tex_packing(nir_builder *b, nir_tex_instr *tex,
       break;
    }
 
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(color),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, color,
                                   color->parent_instr);
 }
 
@@ -1095,7 +1095,7 @@ nir_lower_txs_lod(nir_builder *b, nir_tex_instr *tex)
       minified = nir_vec(b, comp, dest_size);
    }
 
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(minified),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, minified,
                                   minified->parent_instr);
    return true;
 }

src/compiler/nir/nir_lower_texcoord_replace.c

@@ -127,7 +127,7 @@ nir_lower_texcoord_replace_impl(nir_function_impl *impl,
                                          &intrin->dest.ssa);
 
          nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                        nir_src_for_ssa(result),
+                                        result,
                                         result->parent_instr);
       }
    }

src/compiler/nir/nir_lower_vec3_to_vec4.c

@@ -77,7 +77,7 @@ lower_vec3_to_vec4_impl(nir_function_impl *impl, nir_variable_mode modes)
                b.cursor = nir_after_instr(&intrin->instr);
                nir_ssa_def *vec3 = nir_channels(&b, &intrin->dest.ssa, 0x7);
                nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                              nir_src_for_ssa(vec3),
+                                              vec3,
                                               vec3->parent_instr);
                progress = true;
                break;

src/compiler/nir/nir_lower_wpos_center.c

@@ -67,7 +67,7 @@ update_fragcoord(nir_builder *b, nir_intrinsic_instr *intr,
                                nir_imm_float(b, 0.0f)));
    }
 
-   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(wpos),
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, wpos,
                                   wpos->parent_instr);
 }
 

src/compiler/nir/nir_lower_wpos_ytransform.c

@@ -140,7 +140,7 @@ emit_wpos_adjustment(lower_wpos_ytransform_state *state,
                         nir_channel(b, wpos_temp, 3));
 
    nir_ssa_def_rewrite_uses_after(&intr->dest.ssa,
-                                  nir_src_for_ssa(wpos_temp),
+                                  wpos_temp,
                                   wpos_temp->parent_instr);
 }
 
@@ -291,7 +291,7 @@ lower_load_sample_pos(lower_wpos_ytransform_state *state,
                         nir_fmul(b, nir_channel(b, pos, 1), scale));
    nir_ssa_def *flipped_pos = nir_vec2(b, nir_channel(b, pos, 0), flipped_y);
 
-   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(flipped_pos),
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, flipped_pos,
                                   flipped_pos->parent_instr);
 }
 

src/compiler/nir/nir_opt_copy_prop_vars.c

@@ -964,7 +964,7 @@ copy_prop_vars_block(struct copy_prop_var_state *state,
                    * rewrite the vecN itself.
                    */
                   nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                                 nir_src_for_ssa(value.ssa.def[0]),
+                                                 value.ssa.def[0],
                                                  value.ssa.def[0]->parent_instr);
                } else {
                   nir_ssa_def_rewrite_uses(&intrin->dest.ssa,

src/compiler/nir/nir_opt_load_store_vectorize.c

@@ -732,12 +732,12 @@ vectorize_loads(nir_builder *b, struct vectorize_ctx *ctx,
 
    /* update uses */
    if (first == low) {
-      nir_ssa_def_rewrite_uses_after(&low->intrin->dest.ssa, nir_src_for_ssa(low_def),
+      nir_ssa_def_rewrite_uses_after(&low->intrin->dest.ssa, low_def,
                                      high_def->parent_instr);
       nir_ssa_def_rewrite_uses(&high->intrin->dest.ssa, high_def);
    } else {
       nir_ssa_def_rewrite_uses(&low->intrin->dest.ssa, low_def);
-      nir_ssa_def_rewrite_uses_after(&high->intrin->dest.ssa, nir_src_for_ssa(high_def),
+      nir_ssa_def_rewrite_uses_after(&high->intrin->dest.ssa, high_def,
                                      high_def->parent_instr);
    }
 

src/compiler/nir/nir_split_vars.c

@@ -1578,7 +1578,7 @@ shrink_vec_var_access_impl(nir_function_impl *impl,
                nir_ssa_def *vec = nir_vec(&b, vec_srcs, intrin->num_components);
 
                nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                              nir_src_for_ssa(vec),
+                                              vec,
                                               vec->parent_instr);
 
                /* The SSA def is now only used by the swizzle.  It's safe to

src/freedreno/ir3/ir3_nir_lower_tess.c

@@ -731,7 +731,7 @@ lower_tess_eval_block(nir_block *block, nir_builder *b, struct state *state)
 			nir_ssa_def *coord = nir_vec3(b, x, y, z);
 
 			nir_ssa_def_rewrite_uses_after(&intr->dest.ssa,
-					nir_src_for_ssa(coord),
+					coord,
 					b->cursor.instr);
 			break;
 		}

src/freedreno/vulkan/tu_shader.c

@@ -505,7 +505,7 @@ lower_tex_ycbcr(const struct tu_pipeline_layout *layout,
                                                   ycbcr_sampler->ycbcr_range,
                                                   &tex->dest.ssa,
                                                   bpcs);
-   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, nir_src_for_ssa(result),
+   nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, result,
                                   result->parent_instr);
 
    builder->cursor = nir_before_instr(&tex->instr);

src/gallium/drivers/d3d12/d3d12_nir_passes.c

@@ -194,7 +194,7 @@ lower_pos_read(nir_builder *b, struct nir_instr *instr,
    pos = nir_vector_insert_imm(b, pos, depth, 2);
 
    assert(intr->dest.is_ssa);
-   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(pos),
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, pos,
                                   pos->parent_instr);
 }
 

src/gallium/drivers/etnaviv/etnaviv_nir.c

@@ -55,7 +55,7 @@ etna_lower_io(nir_shader *shader, struct etna_shader_variant *v)
                      nir_instr_as_alu(ssa->parent_instr)->op = nir_op_ieq;
 
                   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa,
-                                                 nir_src_for_ssa(ssa),
+                                                 ssa,
                                                  ssa->parent_instr);
                } break;
                case nir_intrinsic_store_deref: {
@@ -197,7 +197,8 @@ etna_lower_alu_impl(nir_function_impl *impl, bool has_new_transcendentals)
 
             nir_instr_insert_after(instr, &mul->instr);
 
-            nir_ssa_def_rewrite_uses_after(ssa, nir_src_for_ssa(&mul->dest.dest.ssa), &mul->instr);
+            nir_ssa_def_rewrite_uses_after(ssa, &mul->dest.dest.ssa,
+                                           &mul->instr);
          }
       }
    }

src/gallium/drivers/vc4/vc4_nir_lower_io.c

@@ -267,7 +267,7 @@ vc4_nir_lower_fs_input(struct vc4_compile *c, nir_builder *b,
 
                 if (result != &intr->dest.ssa) {
                         nir_ssa_def_rewrite_uses_after(&intr->dest.ssa,
-                                                       nir_src_for_ssa(result),
+                                                       result,
                                                        result->parent_instr);
                 }
         }

src/gallium/drivers/zink/nir_lower_dynamic_bo_access.c

@@ -129,7 +129,8 @@ lower_dynamic_bo_access_instr(nir_intrinsic_instr *instr, nir_builder *b)
       /* now use the composite dest in all cases where the original dest (from the dynamic index)
        * was used and remove the dynamically-indexed load_*bo instruction
        */
-      nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, nir_src_for_ssa(new_dest), &instr->instr);
+      nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, new_dest,
+                                     &instr->instr);
    } else
       generate_store_ssbo_ssa_def(b, instr, instr->src[block_idx].ssa, first_idx, last_idx);
    nir_instr_remove(&instr->instr);

src/gallium/drivers/zink/zink_compiler.c

@@ -196,7 +196,8 @@ lower_64bit_vertex_attribs_instr(nir_builder *b, nir_instr *instr, void *data)
          def[3] = nir_vector_extract(b, &intr2->dest.ssa, nir_imm_int(b, 1));
       nir_ssa_def *new_vec = nir_vec(b, def, total_num_components);
       /* use the assembled dvec3/4 for all other uses of the load */
-      nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(new_vec), new_vec->parent_instr);
+      nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, new_vec,
+                                     new_vec->parent_instr);
    }
 
    return true;
@@ -241,7 +242,8 @@ lower_basevertex_instr(nir_builder *b, nir_instr *in, void *data)
                                           nir_imm_int(b, 0),
                                           NULL);
 
-   nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, nir_src_for_ssa(composite), composite->parent_instr);
+   nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, composite,
+                                  composite->parent_instr);
    return true;
 }
 
@@ -513,7 +515,7 @@ lower_baseinstance_instr(nir_builder *b, nir_instr *instr, void *data)
       return false;
    b->cursor = nir_after_instr(instr);
    nir_ssa_def *def = nir_isub(b, &intr->dest.ssa, nir_load_base_instance(b));
-   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(def), def->parent_instr);
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, def, def->parent_instr);
    return true;
 }
 

src/intel/compiler/brw_nir_attribute_workarounds.c

@@ -122,7 +122,7 @@ apply_attr_wa_block(nir_block *block, struct attr_wa_state *state)
                                                : nir_u2f32(b, val);
       }
 
-      nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa, nir_src_for_ssa(val),
+      nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa, val,
                                      val->parent_instr);
       state->impl_progress = true;
    }

src/intel/compiler/brw_nir_clamp_image_1d_2d_array_sizes.c

@@ -138,7 +138,7 @@ brw_nir_clamp_image_1d_2d_array_sizes(nir_shader *shader)
             b.cursor = nir_after_instr(instr);
 
             nir_ssa_def_rewrite_uses_after(image_size,
-                                           nir_src_for_ssa(image_size_replacement),
+                                           image_size_replacement,
                                            image_size_replacement->parent_instr);
 
             function_progress = true;

src/intel/compiler/brw_nir_lower_cs_intrinsics.c

@@ -61,7 +61,7 @@ lower_cs_intrinsics_convert_block(struct lower_intrinsics_state *state,
             intrinsic->dest.ssa.bit_size = 32;
             sysval = nir_u2u64(b, &intrinsic->dest.ssa);
             nir_ssa_def_rewrite_uses_after(&intrinsic->dest.ssa,
-                                           nir_src_for_ssa(sysval),
+                                           sysval,
                                            sysval->parent_instr);
          }
          continue;

src/intel/vulkan/anv_nir_add_base_work_group_id.c

@@ -58,7 +58,7 @@ anv_nir_add_base_work_group_id(nir_shader *shader)
                                            load_base);
 
             nir_ssa_def_rewrite_uses_after(&load_id->dest.ssa,
-                                           nir_src_for_ssa(id),
+                                           id,
                                            id->parent_instr);
             progress = true;
          }

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -1031,7 +1031,7 @@ lower_gen7_tex_swizzle(nir_tex_instr *tex, unsigned plane,
 
    /* Rewrite uses before we insert so we don't rewrite this use */
    nir_ssa_def_rewrite_uses_after(&tex->dest.ssa,
-                                  nir_src_for_ssa(swiz_tex_res),
+                                  swiz_tex_res,
                                   swiz_tex_res->parent_instr);
 }
 

src/microsoft/clc/clc_compiler.c

@@ -369,7 +369,7 @@ clc_lower_64bit_semantics(nir_shader *nir)
                nir_ssa_def *i64 = nir_u2u64(&b, &intrinsic->dest.ssa);
                nir_ssa_def_rewrite_uses_after(
                   &intrinsic->dest.ssa,
-                  nir_src_for_ssa(i64),
+                  i64,
                   i64->parent_instr);
             }
          }

src/panfrost/midgard/nir_fuse_io_16.c

@@ -94,7 +94,8 @@ nir_fuse_io_16(nir_shader *shader)
 
             /* The f2f32(f2fmp(x)) will cancel by opt_algebraic */
             nir_ssa_def *conv = nir_f2f32(&b, &intr->dest.ssa);
-            nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(conv), conv->parent_instr);
+            nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, conv,
+                                           conv->parent_instr);
 
             progress |= true;
          }

src/panfrost/util/pan_lower_64bit_intrin.c

@@ -62,7 +62,8 @@ nir_lower_64bit_intrin_instr(nir_builder *b, nir_instr *instr, void *data)
 
         nir_ssa_def *conv = nir_u2u64(b, &intr->dest.ssa);
 
-        nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, nir_src_for_ssa(conv), conv->parent_instr);
+        nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, conv,
+                                       conv->parent_instr);
 
         return true;
 }

src/panfrost/util/pan_lower_framebuffer.c

@@ -707,8 +707,7 @@ pan_lower_fb_load(nir_shader *shader,
         unpacked = nir_convert_to_bit_size(b, unpacked, src_type, bits);
         unpacked = pan_extend(b, unpacked, nir_dest_num_components(intr->dest));
 
-        nir_src rewritten = nir_src_for_ssa(unpacked);
+        nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, unpacked, &intr->instr);
-        nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, rewritten, &intr->instr);
 }
 
 bool