nir: Take a shader and variable mode in nir_assign_io_var_locations