nir/uub: Use an optional max_samples from drivers for sample counts.

This triggers some unrolling in Fallout 4, GTAV, and Rocky Planet in my
shader-db.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38585>

src/amd/common/nir/ac_nir.c

@@ -128,6 +128,7 @@ void ac_nir_set_options(struct radeon_info *info, bool use_llvm,
    options->max_workgroup_count[0] = UINT32_MAX;
    options->max_workgroup_count[1] = UINT16_MAX;
    options->max_workgroup_count[2] = UINT16_MAX;
+   options->max_samples = 8;
 }
 
 /* Sleep for the given number of clock cycles. */

src/amd/vulkan/radv_physical_device.c

@@ -3060,6 +3060,7 @@ radv_GetPhysicalDeviceMultisamplePropertiesEXT(VkPhysicalDevice physicalDevice,
 {
    VK_FROM_HANDLE(radv_physical_device, pdev, physicalDevice);
 
+   /* Note: update nir_shader_compiler_options.max_samples when changing this. */
    VkSampleCountFlagBits supported_samples = VK_SAMPLE_COUNT_2_BIT | VK_SAMPLE_COUNT_4_BIT | VK_SAMPLE_COUNT_8_BIT;
    if (pdev->info.gfx_level >= GFX10)
       supported_samples |= VK_SAMPLE_COUNT_1_BIT;

src/asahi/compiler/agx_compile.h

@@ -401,6 +401,7 @@ static const nir_shader_compiler_options agx_nir_options = {
                               BITFIELD_BIT(MESA_SHADER_TESS_EVAL) |
                               BITFIELD_BIT(MESA_SHADER_FRAGMENT),
    .support_indirect_outputs = (uint8_t)BITFIELD_MASK(MESA_SHADER_STAGES),
+   .max_samples = 4,
    .lower_fquantize2f16 = true,
    .compact_arrays = true,
    .discard_is_demote = true,

src/asahi/vulkan/hk_physical_device.c

@@ -631,6 +631,7 @@ hk_get_device_properties(const struct agx_device *dev,
                          const struct hk_instance *instance,
                          struct vk_properties *properties)
 {
+   /* Note: update nir_shader_compiler_options.max_samples when changing this. */
    const VkSampleCountFlagBits sample_counts =
       VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_2_BIT | VK_SAMPLE_COUNT_4_BIT;
 

src/broadcom/vulkan/v3dv_device.c

@@ -844,6 +844,7 @@ get_device_properties(const struct v3dv_physical_device *device,
    const float v3d_point_line_granularity = 2.0f / (1 << V3D_COORD_SHIFT);
    const uint32_t max_fb_size = V3D_MAX_IMAGE_DIMENSION;
 
+   /* Note: update nir_shader_compiler_options.max_samples when changing this. */
    const VkSampleCountFlags supported_sample_counts =
       VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_4_BIT;
 

src/compiler/nir/nir_range_analysis.c

@@ -1808,6 +1808,12 @@ get_intrinsic_uub(struct analysis_state *state, struct scalar_query q, uint32_t
          *result = upper_bound;
       break;
    }
+
+   case nir_intrinsic_image_samples:
+      if (state->shader->options->max_samples > 0)
+         *result = state->shader->options->max_samples;
+      break;
+
    default:
       break;
    }
@@ -2111,6 +2117,15 @@ get_alu_uub(struct analysis_state *state, struct scalar_query q, uint32_t *resul
    }
 }
 
+static void
+get_tex_uub(struct analysis_state *state, struct scalar_query q, uint32_t *result, const uint32_t *src)
+{
+   nir_tex_instr *tex = nir_scalar_as_tex(q.scalar);
+
+   if (tex->op == nir_texop_texture_samples && state->shader->options->max_samples > 0)
+      *result = state->shader->options->max_samples;
+}
+
 static void
 get_phi_uub(struct analysis_state *state, struct scalar_query q, uint32_t *result, const uint32_t *src)
 {
@@ -2158,6 +2173,8 @@ process_uub_query(struct analysis_state *state, struct analysis_query *aq, uint3
       get_intrinsic_uub(state, q, result, src);
    else if (nir_scalar_is_alu(q.scalar))
       get_alu_uub(state, q, result, src);
+   else if (nir_def_instr_type(q.scalar.def) == nir_instr_type_tex)
+      get_tex_uub(state, q, result, src);
    else if (nir_def_instr_type(q.scalar.def) == nir_instr_type_phi)
       get_phi_uub(state, q, result, src);
 }

src/compiler/nir/nir_shader_compiler_options.h

@@ -756,6 +756,12 @@ typedef struct nir_shader_compiler_options {
    uint8_t support_indirect_inputs;
    uint8_t support_indirect_outputs;
 
+   /**
+    * If set, the maximum MSAA sample count supported -- can hint loop unrolling
+    * to optimistically unroll a loop doing txf_ms per sample.
+    */
+   uint8_t max_samples;
+
    /**
     * Lower fmulz to `min(abs(a), abs(b)) == 0.0 ? 0.0 : a * b`.
     */

src/freedreno/ir3/ir3_compiler.c

@@ -114,6 +114,7 @@ static const nir_shader_compiler_options ir3_base_options = {
    .force_indirect_unrolling_sampler = true,
    .lower_uniforms_to_ubo = true,
    .max_unroll_iterations = 32,
+   .max_samples = 4,
 
    /* Not actually supported but we want fmulz to be produced and then be
     * lowered with the abs min pattern since we have free abs on min.

src/freedreno/vulkan/tu_device.cc

@@ -875,7 +875,7 @@ static const size_t max_descriptor_set_size = MAX_SET_SIZE / (4 * A6XX_TEX_CONST
 static const VkSampleCountFlags sample_counts =
    VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_2_BIT | VK_SAMPLE_COUNT_4_BIT;
 static const VkSampleCountFlags sample_location_counts =
-   VK_SAMPLE_COUNT_2_BIT | VK_SAMPLE_COUNT_4_BIT;
+   VK_SAMPLE_COUNT_2_BIT | VK_SAMPLE_COUNT_4_BIT; /* Note: update nir_shader_compiler_options.max_samples when changing this. */
 
 static void
 tu_get_physical_device_properties_1_2(struct tu_physical_device *pdevice,

src/gallium/drivers/v3d/v3d_screen.c

@@ -620,6 +620,7 @@ v3d_screen_get_compiler_options(struct pipe_screen *pscreen,
                  * limit register pressure impact.
                  */
                 .max_unroll_iterations = 16,
+                .max_samples = 4,
                 .force_indirect_unrolling_sampler = true,
                 .scalarize_ddx = true,
                 .max_varying_expression_cost = 4,

src/gallium/drivers/vc4/vc4_program.c

@@ -2160,6 +2160,7 @@ static const nir_shader_compiler_options nir_options = {
         .has_texture_scaling = true,
         .lower_mul_high = true,
         .max_unroll_iterations = 32,
+        .max_samples = 4,
         .force_indirect_unrolling = (nir_var_shader_in | nir_var_shader_out | nir_var_function_temp),
         .scalarize_ddx = true,
 };

src/imagination/pco/pco_nir.c

@@ -74,6 +74,7 @@ static const nir_shader_compiler_options nir_options = {
    .lower_helper_invocation = true,
 
    .max_unroll_iterations = 16,
+   .max_samples = 4,
 
    .io_options = nir_io_vectorizer_ignores_types,
 };

src/imagination/vulkan/pvr_physical_device.c

@@ -659,6 +659,7 @@ static bool pvr_physical_device_get_properties(
       .maxFramebufferHeight = 4096U,
       .maxFramebufferLayers = 256U,
 
+      /* Note: update nir_shader_compiler_options.max_samples when changing this. */
       .framebufferColorSampleCounts = VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_4_BIT,
       .framebufferDepthSampleCounts = VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_4_BIT,
       .framebufferStencilSampleCounts = VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_4_BIT,

src/intel/compiler/brw/brw_compiler.c

@@ -204,6 +204,7 @@ brw_compiler_create(void *mem_ctx, const struct intel_device_info *devinfo)
 
    nir_options->lower_int64_options = int64_options;
    nir_options->lower_doubles_options = fp64_options;
+   nir_options->max_samples = devinfo->ver >= 30 ? 8 : 16;
 
    if (compiler->use_tcs_multi_patch) {
       /* TCS MULTI_PATCH mode has multiple patches per subgroup */

src/intel/compiler/elk/elk_nir_options.c

@@ -55,6 +55,7 @@
    .avoid_ternary_with_two_constants = true,                                  \
    .has_pack_32_4x8 = true,                                                   \
    .max_unroll_iterations = 32,                                               \
+   .max_samples = 8,                                                          \
    .force_indirect_unrolling = nir_var_function_temp,                         \
    .divergence_analysis_options =                                             \
       (nir_divergence_single_patch_per_tcs_subgroup |                         \

src/intel/isl/isl.c

@@ -507,6 +507,7 @@ isl_device_get_supported_tilings(const struct isl_device *dev)
  *
  * This function always returns non-zero, as ISL_SAMPLE_COUNT_1_BIT is always
  * supported.
+ * Note: update nir_shader_compiler_options.max_samples when changing this.
  */
 isl_sample_count_mask_t ATTRIBUTE_CONST
 isl_device_get_sample_counts(const struct isl_device *dev)

src/nouveau/compiler/nak/api.rs

@@ -191,6 +191,7 @@ fn nir_options(dev: &nv_device_info) -> nir_shader_compiler_options {
     op.discard_is_demote = true;
 
     op.max_unroll_iterations = 32;
+    op.max_samples = 8;
     op.scalarize_ddx = true;
 
     op

src/nouveau/vulkan/nvk_physical_device.c

@@ -762,6 +762,7 @@ nvk_get_device_properties(const struct nvk_instance *instance,
                           const struct nv_device_info *info,
                           struct vk_properties *properties)
 {
+   /* Note: update nir_shader_compiler_options.max_samples when changing this. */
    const VkSampleCountFlagBits sample_counts = VK_SAMPLE_COUNT_1_BIT |
                                                VK_SAMPLE_COUNT_2_BIT |
                                                VK_SAMPLE_COUNT_4_BIT |

src/panfrost/compiler/bifrost/bifrost_compile.h

@@ -174,6 +174,7 @@ valhal_writes_extended_fifo(uint64_t outputs_written,
       .lower_cs_local_index_to_id = true,                                      \
       .lower_device_index_to_zero = true,                                      \
       .max_unroll_iterations = 32,                                             \
+      .max_samples = 16,                                                       \
       .force_indirect_unrolling =                                              \
          (nir_var_shader_in | nir_var_shader_out | nir_var_function_temp),     \
       .force_indirect_unrolling_sampler = true,                                \

src/panfrost/compiler/midgard/midgard_compile.h

@@ -104,6 +104,7 @@ static const nir_shader_compiler_options midgard_nir_options = {
    .has_cs_global_id = true,
    .lower_cs_local_index_to_id = true,
    .max_unroll_iterations = 32,
+   .max_samples = 16,
    .force_indirect_unrolling =
       (nir_var_shader_in | nir_var_shader_out | nir_var_function_temp),
    .force_indirect_unrolling_sampler = true,

src/panfrost/vulkan/panvk_physical_device.c

@@ -707,6 +707,7 @@ get_image_format_features(struct panvk_physical_device *physical_device,
    return features;
 }
 
+/* Note: update nir_shader_compiler_options.max_samples when changing this. */
 VkSampleCountFlags
 panvk_get_sample_counts(unsigned arch, unsigned max_tib_size,
                         unsigned max_cbuf_atts, unsigned format_size)