From feb2d3e1daabf71e7d63497541ab2f08a0b7c173 Mon Sep 17 00:00:00 2001
From: Daniel Almeida <daniel.almeida@collabora.com>
Date: Mon, 15 Jan 2024 15:35:50 -0300
Subject: [PATCH] nak/sm50: support annotations through  OpAnnotate

Add a new op to annotate the IR. This will help debugging and is only
in effect when NAK_DEBUG=annotate is set.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27158>
---
 src/nouveau/compiler/nak/api.rs      |  2 ++
 src/nouveau/compiler/nak/from_nir.rs | 14 +++++++++
 src/nouveau/compiler/nak/ir.rs       | 47 +++++++++++++++++++++++++---
 3 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/src/nouveau/compiler/nak/api.rs b/src/nouveau/compiler/nak/api.rs
index 865e001f47a..766c76af43c 100644
--- a/src/nouveau/compiler/nak/api.rs
+++ b/src/nouveau/compiler/nak/api.rs
@@ -424,6 +424,8 @@ pub extern "C" fn nak_compile_shader(
         write!(asm, "{}", s).expect("Failed to dump assembly");
     }
 
+    s.remove_annotations();
+
     let code = if nak.sm >= 70 {
         s.encode_sm70()
     } else if nak.sm >= 50 {
diff --git a/src/nouveau/compiler/nak/from_nir.rs b/src/nouveau/compiler/nak/from_nir.rs
index 3cf85f1b17f..538af443621 100644
--- a/src/nouveau/compiler/nak/from_nir.rs
+++ b/src/nouveau/compiler/nak/from_nir.rs
@@ -3,6 +3,8 @@
 
 #![allow(non_upper_case_globals)]
 
+use crate::api::GetDebugFlags;
+use crate::api::DEBUG;
 use crate::cfg::CFGBuilder;
 use crate::ir::*;
 use crate::nir::*;
@@ -2881,6 +2883,18 @@ impl<'a> ShaderFromNir<'a> {
         }
 
         for ni in nb.iter_instr_list() {
+            if DEBUG.annotate() {
+                let annotation = self
+                    .nir_instr_printer
+                    .instr_to_string(ni)
+                    .split_whitespace()
+                    .collect::<Vec<_>>()
+                    .join(" ");
+                b.push_op(OpAnnotate {
+                    annotation: format!("generated by \"{}\"", annotation,),
+                });
+            }
+
             match ni.type_ {
                 nir_instr_type_alu => {
                     self.parse_alu(&mut b, ni.as_alu().unwrap())
diff --git a/src/nouveau/compiler/nak/ir.rs b/src/nouveau/compiler/nak/ir.rs
index d41ef0bd329..445296a3cb5 100644
--- a/src/nouveau/compiler/nak/ir.rs
+++ b/src/nouveau/compiler/nak/ir.rs
@@ -4850,6 +4850,26 @@ impl DisplayOp for OpOutFinal {
 }
 impl_display_for_op!(OpOutFinal);
 
+/// Describes an annotation on an instruction.
+#[repr(C)]
+#[derive(SrcsAsSlice, DstsAsSlice)]
+pub struct OpAnnotate {
+    /// The annotation
+    pub annotation: String,
+}
+
+impl DisplayOp for OpAnnotate {
+    fn fmt_op(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "// {}", self.annotation)
+    }
+}
+
+impl fmt::Display for OpAnnotate {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.fmt_op(f)
+    }
+}
+
 #[derive(DisplayOp, DstsAsSlice, SrcsAsSlice, FromVariants)]
 pub enum Op {
     FAdd(OpFAdd),
@@ -4944,6 +4964,7 @@ pub enum Op {
     FSOut(OpFSOut),
     Out(OpOut),
     OutFinal(OpOutFinal),
+    Annotate(OpAnnotate),
 }
 impl_display_for_op!(Op);
 
@@ -5277,7 +5298,8 @@ impl Instr {
             | Op::Bar(_)
             | Op::FSOut(_)
             | Op::Out(_)
-            | Op::OutFinal(_) => false,
+            | Op::OutFinal(_)
+            | Op::Annotate(_) => false,
             Op::BMov(op) => !op.clear,
             _ => true,
         }
@@ -5394,7 +5416,8 @@ impl Instr {
             | Op::Copy(_)
             | Op::Swap(_)
             | Op::ParCopy(_)
-            | Op::FSOut(_) => {
+            | Op::FSOut(_)
+            | Op::Annotate(_) => {
                 panic!("Not a hardware opcode")
             }
         }
@@ -5649,8 +5672,9 @@ impl fmt::Display for Function {
                 pred_width = max(pred_width, pred.len());
                 dsts_width = max(dsts_width, dsts.len());
                 op_width = max(op_width, op.len());
+                let is_annotation = matches!(i.op, Op::Annotate(_));
 
-                instrs.push((pred, dsts, op, deps));
+                instrs.push((pred, dsts, op, deps, is_annotation));
             }
             blocks.push(instrs);
         }
@@ -5665,9 +5689,11 @@ impl fmt::Display for Function {
             }
             write!(f, "] -> {{\n")?;
 
-            for (pred, dsts, op, deps) in b.drain(..) {
+            for (pred, dsts, op, deps, is_annotation) in b.drain(..) {
                 let eq_sym = if dsts.is_empty() { " " } else { "=" };
-                if deps.is_empty() {
+                if is_annotation {
+                    write!(f, "\n{}\n", op)?;
+                } else if deps.is_empty() {
                     write!(
                         f,
                         "{:<pred_width$} {:<dsts_width$} {} {}\n",
@@ -5898,6 +5924,17 @@ impl Shader {
         }
     }
 
+    /// Remove all annotations, presumably before encoding the shader.
+    pub fn remove_annotations(&mut self) {
+        self.map_instrs(|instr: Box<Instr>, _| -> MappedInstrs {
+            if matches!(instr.op, Op::Annotate(_)) {
+                MappedInstrs::None
+            } else {
+                MappedInstrs::One(instr)
+            }
+        })
+    }
+
     pub fn lower_ineg(&mut self) {
         let sm = self.info.sm;
         self.map_instrs(|mut instr: Box<Instr>, _| -> MappedInstrs {