nir: Add an LOD parameter to image_*_size

The OpenCL image_width/height/depth functions have variants which can
take an LOD parameter.  More importantly, LLVM-SPIRV-Translator always
generates OpImageQuerySizeLod even if the LOD is guaranteed to be zero.
Given that over half the hardware out there has an LOD field for image
size queries (based on a rudimentary scan through their NIR -> whatever
code), we may as well just add the source to the NIR intrinsic.  If this
is ever a problem for anyone, the lowering is pretty trivial.

I've also added asserts to everyone's drivers that should alert them if
they ever see an LOD other than zero.  This will never happen with GL or
Vulkan so there's no need for panic.

Reviewed-by: Kenneth Graunke <kenneth@whitecape.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6396>

src/amd/compiler/aco_instruction_selection.cpp

@@ -6074,6 +6074,7 @@ void visit_image_size(isel_context *ctx, nir_intrinsic_instr *instr)
    }
 
    /* LOD */
+   assert(nir_src_as_uint(instr->src[1]) == 0);
    Temp lod = bld.vop1(aco_opcode::v_mov_b32, bld.def(v1), Operand(0u));
 
    /* Resource */

src/amd/llvm/ac_nir_to_llvm.c

@@ -3144,6 +3144,7 @@ static LLVMValueRef visit_image_size(struct ac_nir_context *ctx,
 		args.dmask = 0xf;
 		args.resource = get_image_descriptor(ctx, instr, dynamic_index, AC_DESC_IMAGE, false);
 		args.opcode = ac_image_get_resinfo;
+		assert(nir_src_as_uint(instr->src[1]) == 0);
 		args.lod = ctx->ac.i32_0;
 		args.attributes = AC_FUNC_ATTR_READNONE;
 

src/broadcom/compiler/nir_to_vir.c

@@ -1780,6 +1780,8 @@ ntq_emit_image_size(struct v3d_compile *c, nir_intrinsic_instr *instr)
         unsigned image_index = nir_src_as_uint(instr->src[0]);
         bool is_array = nir_intrinsic_image_array(instr);
 
+        assert(nir_src_as_uint(instr->src[1]) == 0);
+
         ntq_store_dest(c, &instr->dest, 0,
                        vir_uniform(c, QUNIFORM_IMAGE_WIDTH, image_index));
         if (instr->num_components > 1) {

src/compiler/glsl/glsl_to_nir.cpp

@@ -1338,6 +1338,11 @@ nir_visitor::visit(ir_call *ir)
 
          if (op == nir_intrinsic_image_deref_size ||
              op == nir_intrinsic_image_deref_samples) {
+            /* image_deref_size takes an LOD parameter which is always 0
+             * coming from GLSL.
+             */
+            if (op == nir_intrinsic_image_deref_size)
+               instr->src[1] = nir_src_for_ssa(nir_imm_int(&b, 0));
             nir_builder_instr_insert(&b, &instr->instr);
             break;
          }

src/compiler/nir/nir_intrinsics.py

@@ -406,7 +406,7 @@ image("atomic_xor",  src_comp=[4, 1, 1], dest_comp=1)
 image("atomic_exchange",  src_comp=[4, 1, 1], dest_comp=1)
 image("atomic_comp_swap", src_comp=[4, 1, 1, 1], dest_comp=1)
 image("atomic_fadd",  src_comp=[4, 1, 1], dest_comp=1)
-image("size",    dest_comp=0, flags=[CAN_ELIMINATE, CAN_REORDER])
+image("size",    dest_comp=0, src_comp=[1], flags=[CAN_ELIMINATE, CAN_REORDER])
 image("samples", dest_comp=1, flags=[CAN_ELIMINATE, CAN_REORDER])
 image("atomic_inc_wrap",  src_comp=[4, 1, 1], dest_comp=1)
 image("atomic_dec_wrap",  src_comp=[4, 1, 1], dest_comp=1)

src/compiler/spirv/spirv_to_nir.c

@@ -3090,8 +3090,10 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
 
    intrin->src[0] = nir_src_for_ssa(&image.image->dest.ssa);
 
-   /* ImageQuerySize doesn't take any extra parameters */
-   if (opcode != SpvOpImageQuerySize) {
+   if (opcode == SpvOpImageQuerySize) {
+      /* ImageQuerySize only has an LOD which is currently always 0 */
+      intrin->src[1] = nir_src_for_ssa(nir_imm_int(&b->nb, 0));
+   } else {
       /* The image coordinate is always 4 components but we may not have that
        * many.  Swizzle to compensate.
        */

src/freedreno/ir3/ir3_compiler_nir.c

@@ -1219,6 +1219,7 @@ emit_intrinsic_image_size_tex(struct ir3_context *ctx, nir_intrinsic_instr *intr
 			TYPE_U16 : TYPE_U32;
 
 	info.flags |= flags;
+	assert(nir_src_as_uint(intr->src[1]) == 0);
 	lod = create_immed(b, 0);
 	sam = emit_sam(ctx, OPC_GETSIZE, info, dst_type, 0b1111, lod, NULL);
 

src/gallium/drivers/nouveau/codegen/nv50_ir_from_nir.cpp

@@ -2225,6 +2225,7 @@ Converter::visit(nir_intrinsic_instr *insn)
          break;
       case nir_intrinsic_bindless_image_size:
       case nir_intrinsic_image_size:
+         assert(nir_src_as_uint(insn->src[1]) == 0);
          ty = TYPE_U32;
          bindless = op == nir_intrinsic_bindless_image_size;
          break;

src/gallium/drivers/r600/sfn/sfn_emitssboinstruction.cpp

@@ -557,6 +557,8 @@ bool EmitSSBOInstruction::emit_image_size(const nir_intrinsic_instr *intrin)
    GPRVector dest = vec_from_nir(intrin->dest, nir_dest_num_components(intrin->dest));
    GPRVector src{0,{4,4,4,4}};
 
+   assert(nir_src_as_uint(intrin->src[1]) == 0);
+
    auto const_offset = nir_src_as_const_value(intrin->src[0]);
    auto dyn_offset = PValue();
    int res_id = R600_IMAGE_REAL_RESOURCE_OFFSET;

src/intel/compiler/brw_fs_nir.cpp

@@ -4166,6 +4166,8 @@ fs_visitor::nir_emit_intrinsic(const fs_builder &bld, nir_intrinsic_instr *instr
                             BRW_REGISTER_TYPE_UD);
       image = bld.emit_uniformize(image);
 
+      assert(nir_src_as_uint(instr->src[1]) == 0);
+
       fs_reg srcs[TEX_LOGICAL_NUM_SRCS];
       if (instr->intrinsic == nir_intrinsic_image_size)
          srcs[TEX_LOGICAL_SRC_SURFACE] = image;

src/intel/compiler/brw_nir_lower_image_load_store.c

@@ -653,6 +653,8 @@ lower_image_size_instr(nir_builder *b,
    if (isl_has_matching_typed_storage_image_format(devinfo, image_fmt))
       return false;
 
+   assert(nir_src_as_uint(intrin->src[1]) == 0);
+
    b->cursor = nir_instr_remove(&intrin->instr);
 
    nir_ssa_def *size = load_image_param(b, deref, SIZE);