kk: Enable fragmentStoresAndAtomics

Metal will prematurely discard fragments with side effects even if those
side effects happen before the discard. Work around this by making said
discards "optional".

Reviewed-by: Arcady Goldmints-Orlov <arcady@lunarg.com>
Signed-off-by: Aitor Camacho <aitor@lunarg.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38741>

docs/drivers/kosmickrisp/workarounds.rst

@@ -49,6 +49,24 @@ info on what was updated.
 Workarounds
 ===========
 
+KK_WORKAROUND_5
+---------------
+| macOS version: 26.0.1
+| Metal ticket: Not reported
+| Metal ticket status:
+| CTS test failure: ``dEQP-VK.fragment_operations.early_fragment.discard_no_early_fragment_tests_depth``
+| Comments:
+
+Fragment shaders that have side effects (like writing to a buffer) will be
+prematurely discarded if there is a ``discard_fragment`` that will always
+execute. To work around this, we just make the discard "optional" by moving
+it inside a run time conditional that will always be true (such as is the
+fragment a helper?). This tricks the MSL compiler into not optimizing it into
+a premature discard.
+
+| Log:
+| 2025-12-01: Workaround implemented
+
 KK_WORKAROUND_4
 ---------------
 | macOS version: 26.0.1
@@ -98,8 +116,8 @@ The way to fix this is by changing the conditional to:
 KK_WORKAROUND_2
 ---------------
 | macOS version: 15.4.x
-| Metal ticket: Not reported
-| Metal ticket status:
+| Metal ticket: FB21065475 (@aitor)
+| Metal ticket status: Waiting resolution
 | CTS test crash: ``dEQP-VK.graphicsfuzz.cov-nested-loops-never-change-array-element-one`` and ``dEQP-VK.graphicsfuzz.disc-and-add-in-func-in-loop``
 | Comments:
 

src/kosmickrisp/compiler/msl_nir_lower_common.c

@@ -283,6 +283,32 @@ msl_nir_vs_io_types(nir_shader *nir)
                                      NULL);
 }
 
+static bool
+fake_guard_for_discards(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
+{
+   if (intrin->intrinsic != nir_intrinsic_demote)
+      return false;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+   nir_def *helper = nir_is_helper_invocation(b, 1);
+   nir_demote_if(b, nir_inot(b, helper));
+   nir_instr_remove(&intrin->instr);
+   return true;
+}
+
+bool
+msl_nir_fake_guard_for_discards(struct nir_shader *nir)
+{
+   assert(nir->info.stage == MESA_SHADER_FRAGMENT);
+
+   /* No side effects, no lowering needed */
+   if (!nir->info.writes_memory)
+      return false;
+
+   return nir_shader_intrinsics_pass(nir, fake_guard_for_discards,
+                                     nir_metadata_control_flow, NULL);
+}
+
 static bool
 lower_clip_distance(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 {

src/kosmickrisp/compiler/nir_to_msl.h

@@ -55,4 +55,5 @@ bool msl_ensure_depth_write(nir_shader *nir);
 bool msl_ensure_vertex_position_output(nir_shader *nir);
 bool msl_nir_fs_io_types(nir_shader *nir);
 bool msl_nir_vs_io_types(nir_shader *nir);
+bool msl_nir_fake_guard_for_discards(struct nir_shader *nir);
 void msl_lower_nir_late(nir_shader *nir);

src/kosmickrisp/vulkan/kk_physical_device.c

@@ -155,13 +155,7 @@ kk_get_device_features(
       .depthClamp = true,
       .drawIndirectFirstInstance = true,
       .dualSrcBlend = true,
-      /* TODO_KOSMICKRISP
-       * Enabling fragmentStoresAndAtomics fails the following CTS tests, need
-       * to investigate:
-       * dEQP-VK.fragment_operations.early_fragment.discard_no_early_fragment_tests_depth
-       * dEQP-VK.robustness.image_robustness.bind.notemplate.*i.unroll.nonvolatile.sampled_image.no_fmt_qual.img.samples_1.*d_array.frag
-       */
-      .fragmentStoresAndAtomics = false,
+      .fragmentStoresAndAtomics = true,
       .fullDrawIndexUint32 = true,
       .imageCubeArray = true,
       .independentBlend = true,

src/kosmickrisp/vulkan/kk_shader.c

@@ -66,6 +66,7 @@ kk_get_nir_options(struct vk_physical_device *vk_pdev, mesa_shader_stage stage,
       .lower_doubles_options = (nir_lower_doubles_options)(~0),
       .lower_int64_options =
          nir_lower_ufind_msb64 | nir_lower_subgroup_shuffle64,
+      .io_options = nir_io_mediump_is_32bit,
    };
    return &options;
 }
@@ -417,6 +418,9 @@ kk_lower_fs(struct kk_device *dev, nir_shader *nir,
             nir->info.fs.needs_coarse_quad_helper_invocations)
       NIR_PASS(_, nir, msl_lower_static_sample_mask, 0xFFFFFFFF);
 
+   /* KK_WORKAROUND_5 */
+   if (!(dev->disabled_workarounds & BITFIELD64_BIT(5)))
+      NIR_PASS(_, nir, msl_nir_fake_guard_for_discards);
    /* KK_WORKAROUND_4 */
    if (!(dev->disabled_workarounds & BITFIELD64_BIT(4))) {
       NIR_PASS(_, nir, nir_lower_helper_writes, true);