agx: Add algebraic opt to help with discard lowering

When lowering discards, it will be convenient to generate the pattern:

   (cond ? 255 : 0) ^ 255

Add rules to optimize that to

   (cond ? 0 : 255)

This is not part of the main algebraic optimizer since this lowering happens
late.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23998>

src/asahi/compiler/agx_compile.c

@@ -2411,7 +2411,18 @@ agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key,
     */
    if (key->fs.nr_samples) {
       assert(nir->info.stage == MESA_SHADER_FRAGMENT);
-      NIR_PASS_V(nir, agx_nir_lower_sample_mask, key->fs.nr_samples);
+
+      if (agx_nir_lower_sample_mask(nir, key->fs.nr_samples)) {
+         /* Clean up ixor(bcsel) patterns created from sample mask lowering.
+          * If this succeeds, we'll have expressions to constant fold to get the
+          * benefit. We need to rescalarize after folding constants.
+          */
+         if (agx_nir_opt_ixor_bcsel(nir)) {
+            NIR_PASS_V(nir, nir_opt_constant_folding);
+            NIR_PASS_V(nir, nir_lower_load_const_to_scalar);
+            NIR_PASS_V(nir, nir_opt_dce);
+         }
+      }
    }
 
    /* Must be last since NIR passes can remap driver_location freely */

src/asahi/compiler/agx_nir.h

@@ -9,6 +9,7 @@
 
 struct nir_shader;
 
+bool agx_nir_opt_ixor_bcsel(struct nir_shader *shader);
 bool agx_nir_lower_algebraic_late(struct nir_shader *shader);
 bool agx_nir_fuse_algebraic_late(struct nir_shader *shader);
 

src/asahi/compiler/agx_nir_algebraic.py

@@ -121,6 +121,12 @@ for s in range(1, 5):
         (('ishl', a, s), iaddshl(0, a, s)),
     ]
 
+# Discard lowering generates this pattern, clean it up
+ixor_bcsel = [
+   (('ixor', ('bcsel', a, '#b', '#c'), '#d'),
+    ('bcsel', a, ('ixor', b, d), ('ixor', c, d))),
+]
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('-p', '--import-path', required=True)
@@ -137,6 +143,8 @@ def run():
                                       lower_sm5_shift + lower_pack).render())
     print(nir_algebraic.AlgebraicPass("agx_nir_fuse_algebraic_late",
                                       fuse_imad).render())
+    print(nir_algebraic.AlgebraicPass("agx_nir_opt_ixor_bcsel",
+                                      ixor_bcsel).render())
 
 
 if __name__ == '__main__':