nir: Use nir_shader::stage rather than passing it around.