nak: Add geometry shader support

Tested on SM75 and SM86.

Signed-off-by: Mary Guillemard <mary.guillemard@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24998>

src/nouveau/compiler/meson.build

@@ -17,6 +17,7 @@ libnak_c_files = files(
   'nak_nir.c',
   'nak_nir_lower_tex.c',
   'nak_nir_lower_vtg_io.c',
+  'nak_nir_lower_gs_intrinsics.c',
 )
 
 libnak_rs_files = files(

src/nouveau/compiler/nak.rs

@@ -19,6 +19,7 @@ mod nak_lower_par_copies;
 mod nak_opt_copy_prop;
 mod nak_opt_dce;
 mod nak_opt_lop;
+mod nak_opt_out;
 mod nak_repair_ssa;
 mod nak_sph;
 mod nak_spill_values;
@@ -251,6 +252,11 @@ pub extern "C" fn nak_compile_shader(
         eprintln!("NAK IR after dce:\n{}", &s);
     }
 
+    s.opt_out();
+    if DEBUG.print() {
+        eprintln!("NAK IR:\n{}", &s);
+    }
+
     s.legalize();
     if DEBUG.print() {
         eprintln!("NAK IR after legalize:\n{}", &s);

src/nouveau/compiler/nak_encode_sm75.rs

@@ -1739,6 +1739,35 @@ impl SM75Instr {
         self.set_field(72..80, op.idx);
     }
 
+    fn encode_out(&mut self, op: &OpOut) {
+        self.encode_alu(
+            0x124,
+            Some(op.dst),
+            ALUSrc::from_src(&op.handle),
+            ALUSrc::from_src(&op.stream),
+            ALUSrc::None,
+        );
+
+        self.set_field(
+            78..80,
+            match op.out_type {
+                OutType::Emit => 1_u8,
+                OutType::Cut => 2_u8,
+                OutType::EmitThenCut => 3_u8,
+            },
+        );
+    }
+
+    fn encode_out_final(&mut self, op: &OpOutFinal) {
+        self.encode_alu(
+            0x124,
+            Some(Dst::None),
+            ALUSrc::from_src(&op.handle),
+            ALUSrc::from_src(&Src::new_zero()),
+            ALUSrc::None,
+        );
+    }
+
     pub fn encode(
         instr: &Instr,
         sm: u8,
@@ -1814,6 +1843,8 @@ impl SM75Instr {
             Op::Nop(op) => si.encode_nop(&op),
             Op::PixLd(op) => si.encode_pixld(&op),
             Op::S2R(op) => si.encode_s2r(&op),
+            Op::Out(op) => si.encode_out(&op),
+            Op::OutFinal(op) => si.encode_out_final(&op),
             _ => panic!("Unhandled instruction"),
         }
 

src/nouveau/compiler/nak_from_nir.rs

@@ -42,13 +42,14 @@ fn init_info_from_nir(nir: &nir_shader, sm: u8) -> ShaderInfo {
             MESA_SHADER_FRAGMENT => ShaderStageInfo::Fragment,
             MESA_SHADER_GEOMETRY => {
                 let info_gs = unsafe { &nir.info.__bindgen_anon_1.gs };
-                let output_topology = match info_gs.input_primitive {
+                let output_topology = match info_gs.output_primitive {
                     MESA_PRIM_POINTS => OutputTopology::PointList,
-                    MESA_PRIM_TRIANGLES | MESA_PRIM_LINE_STRIP => {
-                        OutputTopology::LineStrip
-                    }
+                    MESA_PRIM_LINE_STRIP => OutputTopology::LineStrip,
                     MESA_PRIM_TRIANGLE_STRIP => OutputTopology::TriangleStrip,
-                    _ => panic!("Invalid GS input primitive"),
+                    _ => panic!(
+                        "Invalid GS input primitive {}",
+                        info_gs.input_primitive
+                    ),
                 };
 
                 ShaderStageInfo::Geometry(GeometryShaderInfo {
@@ -1881,6 +1882,36 @@ impl<'a> ShaderFromNir<'a> {
                     access: access,
                 });
             }
+            nir_intrinsic_emit_vertex_nv | nir_intrinsic_end_primitive_nv => {
+                assert!(intrin.def.bit_size() == 32);
+                assert!(intrin.def.num_components() == 1);
+
+                let dst = b.alloc_ssa(RegFile::GPR, 1);
+                let handle = self.get_src(&srcs[0]);
+                let stream_id = intrin.stream_id();
+
+                b.push_op(OpOut {
+                    dst: dst.into(),
+                    handle: handle,
+                    stream: stream_id.into(),
+                    out_type: if intrin.intrinsic
+                        == nir_intrinsic_emit_vertex_nv
+                    {
+                        OutType::Emit
+                    } else {
+                        OutType::Cut
+                    },
+                });
+                self.set_dst(&intrin.def, dst);
+            }
+
+            nir_intrinsic_final_primitive_nv => {
+                let handle = self.get_src(&srcs[0]);
+
+                if self.info.sm >= 70 {
+                    b.push_op(OpOutFinal { handle: handle });
+                }
+            }
             _ => panic!(
                 "Unsupported intrinsic instruction: {}",
                 intrin.info().name()
@@ -2219,8 +2250,8 @@ impl<'a> ShaderFromNir<'a> {
                     io.mark_attrs_written(tc..(tc + 8));
                 }
                 _ => panic!("Tessellation must have ShaderIoInfo::Vtg"),
-            }
-            _ => ()
+            },
+            _ => (),
         }
 
         Shader {

src/nouveau/compiler/nak_ir.rs

@@ -896,7 +896,7 @@ impl fmt::Display for SrcRef {
     }
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq)]
 pub enum SrcMod {
     None,
     FAbs,
@@ -981,7 +981,7 @@ pub enum SrcType {
     Pred,
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq)]
 pub struct Src {
     pub src_ref: SrcRef,
     pub src_mod: SrcMod,
@@ -4000,6 +4000,60 @@ impl fmt::Display for OpFSOut {
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub enum OutType {
+    Emit,
+    Cut,
+    EmitThenCut,
+}
+
+impl fmt::Display for OutType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            OutType::Emit => write!(f, "EMIT"),
+            OutType::Cut => write!(f, "CUT"),
+            OutType::EmitThenCut => write!(f, "EMIT_THEN_CUT"),
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(SrcsAsSlice, DstsAsSlice)]
+pub struct OpOut {
+    pub dst: Dst,
+
+    #[src_type(SSA)]
+    pub handle: Src,
+
+    #[src_type(ALU)]
+    pub stream: Src,
+
+    pub out_type: OutType,
+}
+
+impl fmt::Display for OpOut {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "OUT.{} {} {{ {}, {} }}",
+            self.out_type, self.dst, self.handle, self.stream
+        )
+    }
+}
+
+#[repr(C)]
+#[derive(SrcsAsSlice, DstsAsSlice)]
+pub struct OpOutFinal {
+    #[src_type(SSA)]
+    pub handle: Src,
+}
+
+impl fmt::Display for OpOutFinal {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "OUT.FINAL {{ {} }}", self.handle)
+    }
+}
+
 #[derive(Display, DstsAsSlice, SrcsAsSlice, FromVariants)]
 pub enum Op {
     FAdd(OpFAdd),
@@ -4072,6 +4126,8 @@ pub enum Op {
     Swap(OpSwap),
     ParCopy(OpParCopy),
     FSOut(OpFSOut),
+    Out(OpOut),
+    OutFinal(OpOutFinal),
 }
 
 #[derive(Clone, Copy, Eq, Hash, PartialEq)]
@@ -4411,7 +4467,9 @@ impl Instr {
             | Op::Exit(_)
             | Op::WarpSync(_)
             | Op::Bar(_)
-            | Op::FSOut(_) => false,
+            | Op::FSOut(_)
+            | Op::Out(_)
+            | Op::OutFinal(_) => false,
             Op::BMov(op) => !op.clear,
             Op::Nop(op) => op.label.is_none(),
             _ => true,
@@ -4486,6 +4544,9 @@ impl Instr {
             Op::Bra(_) | Op::Exit(_) => true,
             Op::WarpSync(_) => false,
 
+            // Geometry ops
+            Op::Out(_) | Op::OutFinal(_) => false,
+
             // Miscellaneous ops
             Op::Bar(_)
             | Op::CS2R(_)

src/nouveau/compiler/nak_legalize.rs

@@ -252,6 +252,13 @@ fn legalize_instr(b: &mut impl SSABuilder, instr: &mut Instr) {
             copy_src_if_cbuf(b, &mut op.lane, RegFile::GPR);
             copy_src_if_cbuf(b, &mut op.c, RegFile::GPR);
         }
+        Op::Out(op) => {
+            copy_src_if_not_reg(b, &mut op.handle, RegFile::GPR);
+            copy_src_if_cbuf(b, &mut op.stream, RegFile::GPR);
+        }
+        Op::OutFinal(op) => {
+            copy_src_if_not_reg(b, &mut op.handle, RegFile::GPR);
+        }
         Op::Ldc(_) => (), // Nothing to do
         Op::BMov(_) | Op::BSSy(_) | Op::BSync(_) => (), // Nothing to do
         Op::Copy(_) => (), // Nothing to do

src/nouveau/compiler/nak_nir.c

@@ -325,7 +325,8 @@ nak_nir_lower_system_value_instr(nir_builder *b, nir_instr *instr, void *data)
 
    case nir_intrinsic_load_primitive_id: {
       assert(b->shader->info.stage == MESA_SHADER_TESS_CTRL ||
-             b->shader->info.stage == MESA_SHADER_TESS_EVAL);
+             b->shader->info.stage == MESA_SHADER_TESS_EVAL ||
+             b->shader->info.stage == MESA_SHADER_GEOMETRY);
       val = nir_load_per_vertex_input(b, 1, 32, nir_imm_int(b, 0),
                                       nir_imm_int(b, 0),
                                       .base = NAK_ATTR_PRIMITIVE_ID,
@@ -749,7 +750,6 @@ nak_postprocess_nir(nir_shader *nir,
 
    case MESA_SHADER_TESS_CTRL:
    case MESA_SHADER_TESS_EVAL:
-   case MESA_SHADER_GEOMETRY:
       OPT(nir, nak_nir_lower_varyings, nir_var_shader_in | nir_var_shader_out);
       OPT(nir, nir_opt_constant_folding);
       OPT(nir, nak_nir_lower_vtg_io, nak);
@@ -760,6 +760,13 @@ nak_postprocess_nir(nir_shader *nir,
       OPT(nir, nak_nir_lower_fs_outputs);
       break;
 
+   case MESA_SHADER_GEOMETRY:
+      OPT(nir, nak_nir_lower_varyings, nir_var_shader_in | nir_var_shader_out);
+      OPT(nir, nir_opt_constant_folding);
+      OPT(nir, nak_nir_lower_vtg_io, nak);
+      OPT(nir, nak_nir_lower_gs_intrinsics);
+      break;
+
    case MESA_SHADER_COMPUTE:
    case MESA_SHADER_KERNEL:
       break;

src/nouveau/compiler/nak_nir_lower_gs_intrinsics.c

@@ -0,0 +1,125 @@
+/*
+ * Copyright © 2023 Collabora, Ltd.
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nak_private.h"
+#include "nir_builder.h"
+#include "nir_format_convert.h"
+
+
+struct state {
+   nir_builder *builder;
+   nir_variable *handle_var;
+   bool progress;
+};
+
+static void
+rewrite_emit_vertex(nir_intrinsic_instr *intr, struct state *state)
+{
+   nir_builder *b = state->builder;
+   const unsigned stream = nir_intrinsic_stream_id(intr);
+
+   b->cursor = nir_before_instr(&intr->instr);
+
+   nir_def *gs_handle = nir_load_var(b, state->handle_var);
+   nir_def *gs_handle_out = nir_emit_vertex_nv(b, 32, gs_handle, stream);
+
+   nir_store_var(b, state->handle_var, gs_handle_out, 0x1);
+
+   nir_instr_remove(&intr->instr);
+
+   state->progress = true;
+}
+
+static void
+rewrite_end_primitive(nir_intrinsic_instr *intr, struct state *state)
+{
+   nir_builder *b = state->builder;
+   const unsigned stream = nir_intrinsic_stream_id(intr);
+
+   b->cursor = nir_before_instr(&intr->instr);
+
+   nir_def *gs_handle = nir_load_var(b, state->handle_var);
+   nir_def *gs_handle_out = nir_end_primitive_nv(b, 32, gs_handle, stream);
+
+   nir_store_var(b, state->handle_var, gs_handle_out, 0x1);
+
+   nir_instr_remove(&intr->instr);
+
+   state->progress = true;
+}
+
+static void
+rewrite_ast_nv(nir_intrinsic_instr *intr, struct state *state)
+{
+   nir_builder *b = state->builder;
+
+   b->cursor = nir_before_instr(&intr->instr);
+
+   nir_src *vtx = &intr->src[1];
+   nir_def *gs_handle = nir_load_var(b, state->handle_var);
+   nir_src_rewrite(vtx, gs_handle);
+
+   state->progress = true;
+}
+
+static void
+append_final_primitive_nv(nir_block *end_block, struct state *state)
+{
+   nir_builder *b = state->builder;
+
+   set_foreach(end_block->predecessors, entry) {
+      nir_block *pred = (nir_block *)entry->key;
+      b->cursor = nir_after_block_before_jump(pred);
+
+      nir_def *gs_handle = nir_load_var(b, state->handle_var);
+      nir_final_primitive_nv(b, gs_handle);
+
+      state->progress = true;
+   }
+}
+
+bool
+nak_nir_lower_gs_intrinsics(nir_shader *nir)
+{
+   struct state state;
+
+   state.progress = false;
+   nir_function_impl *impl = nir_shader_get_entrypoint(nir);
+   nir_builder builder = nir_builder_at(nir_before_impl(impl));
+
+   state.builder = &builder;
+
+   state.handle_var = nir_local_variable_create(impl, glsl_uint_type(), "gs_handle");
+   nir_store_var(&builder, state.handle_var, nir_imm_int(&builder, 0), 0x1);
+
+   nir_foreach_block_safe(block, impl) {
+      nir_foreach_instr_safe(instr, block) {
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+
+         switch (intr->intrinsic) {
+         case nir_intrinsic_emit_vertex:
+            rewrite_emit_vertex(intr, &state);
+            break;
+         case nir_intrinsic_end_primitive:
+            rewrite_end_primitive(intr, &state);
+            break;
+         case nir_intrinsic_ast_nv:
+            rewrite_ast_nv(intr, &state);
+            break;
+         default:
+            break;
+         }
+      }
+   }
+
+   append_final_primitive_nv(impl->end_block, &state);
+
+   nir_metadata_preserve(impl, nir_metadata_none);
+
+   return state.progress;
+}

src/nouveau/compiler/nak_opt_out.rs

@@ -0,0 +1,65 @@
+/*
+ * Copyright © 2023 Collabora, Ltd.
+ * SPDX-License-Identifier: MIT
+ */
+
+use crate::nak_ir::*;
+
+use std::collections::HashMap;
+use std::slice;
+
+struct OutPass;
+
+impl Shader {
+    pub fn opt_out(&mut self) {
+        if let ShaderStageInfo::Geometry(_) = self.info.stage {
+            for f in &mut self.functions {
+                for b in &mut f.blocks {
+                    let mut instrs = Vec::new();
+
+                    {
+                        let mut drain = b.instrs.drain(..);
+
+                        while let Some(instr) = drain.next() {
+                            match instr.op {
+                                Op::Out(op) if op.out_type == OutType::Emit => {
+                                    let next_op_opt =
+                                        drain.next().map(|x| x.op);
+
+                                    match next_op_opt {
+                                        Some(Op::Out(next_op))
+                                            if next_op.out_type
+                                                == OutType::Cut
+                                                && op.stream
+                                                    == next_op.stream =>
+                                        {
+                                            instrs.push(Instr::new_boxed(
+                                                OpOut {
+                                                    dst: next_op.dst.clone(),
+                                                    handle: op.handle,
+                                                    stream: op.stream,
+                                                    out_type:
+                                                        OutType::EmitThenCut,
+                                                },
+                                            ));
+                                        }
+                                        Some(next_op) => {
+                                            instrs.push(Instr::new_boxed(op));
+                                            instrs.push(Instr::new_boxed(
+                                                next_op,
+                                            ));
+                                        }
+                                        None => {}
+                                    }
+                                }
+                                _ => instrs.push(instr),
+                            }
+                        }
+                    }
+
+                    b.instrs = instrs;
+                }
+            }
+        }
+    }
+}

src/nouveau/compiler/nak_private.h

@@ -132,6 +132,7 @@ struct nak_nir_tex_flags {
 };
 
 bool nak_nir_lower_tex(nir_shader *nir, const struct nak_compiler *nak);
+bool nak_nir_lower_gs_intrinsics(nir_shader *shader);
 
 struct nak_nir_attr_io_flags {
    bool output : 1;