From 41722c61376eacc346fd459108d4e9697834f2fd Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Thu, 18 Jan 2024 16:49:01 -0600
Subject: [PATCH] nak: Add support for imad on Volta+ and enable it in simple
 cases

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27159>
---
 src/nouveau/compiler/nak/from_nir.rs      | 10 ++++++++++
 src/nouveau/compiler/nak_nir_algebraic.py |  8 +++++++-
 src/nouveau/compiler/nak_private.h        | 15 +++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/nouveau/compiler/nak/from_nir.rs b/src/nouveau/compiler/nak/from_nir.rs
index 2603b95a20e..d76f167f1e7 100644
--- a/src/nouveau/compiler/nak/from_nir.rs
+++ b/src/nouveau/compiler/nak/from_nir.rs
@@ -1091,6 +1091,16 @@ impl<'a> ShaderFromNir<'a> {
                     b.isetp(cmp_type, cmp_op, x.into(), y.into())
                 }
             }
+            nir_op_imad => {
+                assert!(alu.def.bit_size() == 32);
+                let dst = b.alloc_ssa(RegFile::GPR, 1);
+                b.push_op(OpIMad {
+                    dst: dst.into(),
+                    srcs: [srcs[0], srcs[1], srcs[2]],
+                    signed: false,
+                });
+                dst
+            }
             nir_op_imax | nir_op_imin | nir_op_umax | nir_op_umin => {
                 let (tp, min) = match alu.op {
                     nir_op_imax => (IntCmpType::I32, SrcRef::False),
diff --git a/src/nouveau/compiler/nak_nir_algebraic.py b/src/nouveau/compiler/nak_nir_algebraic.py
index 87a2225b038..e5ecb53c20e 100644
--- a/src/nouveau/compiler/nak_nir_algebraic.py
+++ b/src/nouveau/compiler/nak_nir_algebraic.py
@@ -26,6 +26,7 @@ import sys
 
 a = 'a'
 b = 'b'
+c = 'c'
 
 # common conditions to improve readability
 volta = 'nak->sm >= 70 && nak->sm < 75'
@@ -38,6 +39,11 @@ algebraic_lowering = [
     (('umax', 'a', 'b'), ('bcsel', ('ult', a, b), b, a), volta),
 ]
 
+late_optimizations = [
+    (('iadd@32', ('imul(nak_is_only_used_by_iadd)', a, b), c),
+     ('imad', a, b, c), 'nak->sm >= 70'),
+]
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--out', required=True, help='Output file.')
@@ -52,7 +58,7 @@ def main():
             f.write('#include "nak_private.h"')
             f.write(nir_algebraic.AlgebraicPass(
                 "nak_nir_lower_algebraic_late",
-                algebraic_lowering,
+                algebraic_lowering + late_optimizations,
                 [
                     ("const struct nak_compiler *", "nak"),
                 ]).render())
diff --git a/src/nouveau/compiler/nak_private.h b/src/nouveau/compiler/nak_private.h
index b26b1279589..c1877ee55b2 100644
--- a/src/nouveau/compiler/nak_private.h
+++ b/src/nouveau/compiler/nak_private.h
@@ -205,6 +205,21 @@ enum nak_fs_out {
 
 bool nak_nir_add_barriers(nir_shader *nir, const struct nak_compiler *nak);
 
+static inline bool
+nak_is_only_used_by_iadd(const nir_alu_instr *instr)
+{
+   nir_foreach_use(src, &instr->def) {
+      nir_instr *use = nir_src_parent_instr(src);
+      if (use->type != nir_instr_type_alu)
+         return false;
+
+      if (nir_instr_as_alu(use)->op != nir_op_iadd)
+         return false;
+   }
+
+   return true;
+}
+
 struct nak_memstream {
    FILE *stream;
    char *buffer;