nir: Take a nir_shader and variable mode in assign_var_locations