⚠️ VeridianOS Kernel Documentation - This is low-level kernel code. All functions are unsafe unless explicitly marked otherwise. no_std

veridian_kernel/graphics/
shader.rs

1//! Shader Compiler and Executor
2//!
3//! Provides a software shader pipeline using TGSI-like instructions. Shaders
4//! are compiled from high-level descriptions into instruction lists, then
5//! executed per-pixel by the software rasteriser.
6//!
7//! All arithmetic uses integer or 16.16 fixed-point math (no FPU required).
8
9#![allow(dead_code)]
10
11use alloc::{collections::BTreeMap, string::String, vec::Vec};
12
13// ---------------------------------------------------------------------------
14// Fixed-point helpers (16.16)
15// ---------------------------------------------------------------------------
16
17/// 16.16 fixed-point shift.
18const FP_SHIFT: i32 = 16;
19
20/// 1.0 in 16.16 fixed-point.
21const FP_ONE: i32 = 1 << FP_SHIFT;
22
23/// Multiply two 16.16 fixed-point values.
24fn fp_mul(a: i32, b: i32) -> i32 {
25    ((a as i64 * b as i64) >> FP_SHIFT) as i32
26}
27
28/// Integer to 16.16 fixed-point.
29fn fp_from_int(v: i32) -> i32 {
30    v << FP_SHIFT
31}
32
33/// 16.16 fixed-point to integer (truncate).
34fn fp_to_int(v: i32) -> i32 {
35    v >> FP_SHIFT
36}
37
38/// Reciprocal: 1/x in 16.16 fixed-point.
39fn fp_rcp(x: i32) -> i32 {
40    if x == 0 {
41        return 0;
42    }
43    ((FP_ONE as i64 * FP_ONE as i64) / x as i64) as i32
44}
45
46// ---------------------------------------------------------------------------
47// Shader types
48// ---------------------------------------------------------------------------
49
50/// Type of shader in the pipeline.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum ShaderType {
53    /// Processes vertex positions and attributes.
54    Vertex,
55    /// Computes pixel (fragment) colour.
56    Fragment,
57}
58
59/// A uniform value passed to shader programs.
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum UniformValue {
62    /// Single integer.
63    Int(i32),
64    /// 2-component vector.
65    Vec2(i32, i32),
66    /// 4-component vector.
67    Vec4(i32, i32, i32, i32),
68    /// 4x4 matrix (column-major, 16.16 fixed-point).
69    Mat4([i32; 16]),
70}
71
72// ---------------------------------------------------------------------------
73// TGSI-like instructions
74// ---------------------------------------------------------------------------
75
76/// Source operand for a TGSI instruction.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum SrcOperand {
79    /// Register index.
80    Reg(u8),
81    /// Uniform slot.
82    Uniform(u8),
83    /// Immediate 16.16 fixed-point value.
84    Immediate(i32),
85}
86
87impl Default for SrcOperand {
88    fn default() -> Self {
89        Self::Immediate(0)
90    }
91}
92
93/// TGSI-like instruction set.
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub enum TgsiInstruction {
96    /// Move: dst = src
97    MOV { dst: u8, src: SrcOperand },
98    /// Add: dst = src0 + src1
99    ADD {
100        dst: u8,
101        src0: SrcOperand,
102        src1: SrcOperand,
103    },
104    /// Multiply: dst = src0 * src1 (16.16 fp_mul)
105    MUL {
106        dst: u8,
107        src0: SrcOperand,
108        src1: SrcOperand,
109    },
110    /// Multiply-add: dst = src0 * src1 + src2
111    MAD {
112        dst: u8,
113        src0: SrcOperand,
114        src1: SrcOperand,
115        src2: SrcOperand,
116    },
117    /// 3-component dot product
118    DP3 {
119        dst: u8,
120        src0: SrcOperand,
121        src1: SrcOperand,
122    },
123    /// 4-component dot product
124    DP4 {
125        dst: u8,
126        src0: SrcOperand,
127        src1: SrcOperand,
128    },
129    /// Texture sample (stub — reads from texture slot)
130    TEX { dst: u8, coord: SrcOperand },
131    /// Sample from texture (alias for TEX with sampler index)
132    SAMPLE {
133        dst: u8,
134        coord: SrcOperand,
135        sampler: u8,
136    },
137    /// Reciprocal: dst = 1.0 / src
138    RCP { dst: u8, src: SrcOperand },
139    /// Reciprocal square root (integer approximation)
140    RSQ { dst: u8, src: SrcOperand },
141}
142
143// ---------------------------------------------------------------------------
144// Shader program
145// ---------------------------------------------------------------------------
146
147/// A compiled shader program containing instructions and uniform bindings.
148#[derive(Debug, Clone)]
149pub struct ShaderProgram {
150    /// Type of shader.
151    pub shader_type: ShaderType,
152    /// Compiled instruction list.
153    pub instructions: Vec<TgsiInstruction>,
154    /// Named uniform bindings (name -> slot index).
155    pub uniforms: BTreeMap<String, u8>,
156    /// Program label (for debugging).
157    pub label: String,
158}
159
160impl ShaderProgram {
161    /// Create an empty shader program.
162    pub fn new(shader_type: ShaderType, label: &str) -> Self {
163        Self {
164            shader_type,
165            instructions: Vec::new(),
166            uniforms: BTreeMap::new(),
167            label: String::from(label),
168        }
169    }
170
171    /// Add an instruction.
172    pub fn push(&mut self, instr: TgsiInstruction) {
173        self.instructions.push(instr);
174    }
175
176    /// Bind a uniform name to a slot.
177    pub fn bind_uniform(&mut self, name: &str, slot: u8) {
178        self.uniforms.insert(String::from(name), slot);
179    }
180
181    /// Number of instructions.
182    pub fn instruction_count(&self) -> usize {
183        self.instructions.len()
184    }
185}
186
187// ---------------------------------------------------------------------------
188// Shader compiler
189// ---------------------------------------------------------------------------
190
191/// High-level shader description that gets compiled to TGSI instructions.
192#[derive(Debug, Clone, PartialEq, Eq)]
193pub enum ShaderOp {
194    /// Set output colour to a constant.
195    SetColor { r: u8, g: u8, b: u8, a: u8 },
196    /// Multiply output by a uniform colour.
197    TintByUniform { uniform_name: String },
198    /// Sample a texture at the fragment coordinate.
199    SampleTexture { sampler: u8 },
200    /// Apply a simple gradient (top-to-bottom).
201    VerticalGradient { top: u32, bottom: u32 },
202    /// Passthrough (identity — copy input to output).
203    Passthrough,
204}
205
206/// Compiles high-level shader descriptions to TGSI instruction lists.
207pub struct ShaderCompiler;
208
209impl ShaderCompiler {
210    /// Compile a fragment shader from a list of operations.
211    pub fn compile_fragment(label: &str, ops: &[ShaderOp]) -> ShaderProgram {
212        let mut prog = ShaderProgram::new(ShaderType::Fragment, label);
213
214        for op in ops {
215            match op {
216                ShaderOp::SetColor { r, g, b, a } => {
217                    prog.push(TgsiInstruction::MOV {
218                        dst: 0,
219                        src: SrcOperand::Immediate(fp_from_int(*r as i32)),
220                    });
221                    prog.push(TgsiInstruction::MOV {
222                        dst: 1,
223                        src: SrcOperand::Immediate(fp_from_int(*g as i32)),
224                    });
225                    prog.push(TgsiInstruction::MOV {
226                        dst: 2,
227                        src: SrcOperand::Immediate(fp_from_int(*b as i32)),
228                    });
229                    prog.push(TgsiInstruction::MOV {
230                        dst: 3,
231                        src: SrcOperand::Immediate(fp_from_int(*a as i32)),
232                    });
233                }
234                ShaderOp::TintByUniform { uniform_name } => {
235                    let slot = prog.uniforms.len() as u8;
236                    prog.bind_uniform(uniform_name, slot);
237                    prog.push(TgsiInstruction::MUL {
238                        dst: 0,
239                        src0: SrcOperand::Reg(0),
240                        src1: SrcOperand::Uniform(slot),
241                    });
242                }
243                ShaderOp::SampleTexture { sampler } => {
244                    prog.push(TgsiInstruction::SAMPLE {
245                        dst: 0,
246                        coord: SrcOperand::Reg(4), // texture coord register
247                        sampler: *sampler,
248                    });
249                }
250                ShaderOp::VerticalGradient { top, bottom } => {
251                    let tr = fp_from_int(((*top >> 16) & 0xFF) as i32);
252                    let br = fp_from_int(((*bottom >> 16) & 0xFF) as i32);
253                    // Interpolate: result = top + (bottom - top) * t
254                    prog.push(TgsiInstruction::MOV {
255                        dst: 0,
256                        src: SrcOperand::Immediate(tr),
257                    });
258                    prog.push(TgsiInstruction::MOV {
259                        dst: 5,
260                        src: SrcOperand::Immediate(br - tr),
261                    });
262                    prog.push(TgsiInstruction::MAD {
263                        dst: 0,
264                        src0: SrcOperand::Reg(5),
265                        src1: SrcOperand::Reg(6), // t parameter
266                        src2: SrcOperand::Reg(0),
267                    });
268                }
269                ShaderOp::Passthrough => {
270                    // No-op: input registers already hold output values
271                }
272            }
273        }
274
275        prog
276    }
277
278    /// Compile a simple vertex passthrough shader.
279    pub fn compile_vertex_passthrough(label: &str) -> ShaderProgram {
280        let mut prog = ShaderProgram::new(ShaderType::Vertex, label);
281        // Copy position register to output
282        prog.push(TgsiInstruction::MOV {
283            dst: 0,
284            src: SrcOperand::Reg(0),
285        });
286        prog
287    }
288}
289
290// ---------------------------------------------------------------------------
291// Shader executor
292// ---------------------------------------------------------------------------
293
294/// Maximum number of registers in the virtual register file.
295const MAX_REGISTERS: usize = 16;
296
297/// Executes shader programs on pixel data using software rasterisation.
298pub struct ShaderExecutor {
299    /// Register file for the current invocation.
300    registers: [i32; MAX_REGISTERS],
301    /// Uniform values bound for the current program.
302    uniform_values: Vec<UniformValue>,
303}
304
305impl Default for ShaderExecutor {
306    fn default() -> Self {
307        Self::new()
308    }
309}
310
311impl ShaderExecutor {
312    /// Create a new executor with zeroed registers.
313    pub fn new() -> Self {
314        Self {
315            registers: [0; MAX_REGISTERS],
316            uniform_values: Vec::new(),
317        }
318    }
319
320    /// Set a uniform value for the next execution.
321    pub fn set_uniform(&mut self, slot: usize, value: UniformValue) {
322        while self.uniform_values.len() <= slot {
323            self.uniform_values.push(UniformValue::Int(0));
324        }
325        self.uniform_values[slot] = value;
326    }
327
328    /// Set an input register value.
329    pub fn set_register(&mut self, reg: u8, value: i32) {
330        if (reg as usize) < MAX_REGISTERS {
331            self.registers[reg as usize] = value;
332        }
333    }
334
335    /// Read a register value.
336    pub fn get_register(&self, reg: u8) -> i32 {
337        if (reg as usize) < MAX_REGISTERS {
338            self.registers[reg as usize]
339        } else {
340            0
341        }
342    }
343
344    /// Read a source operand value.
345    fn read_src(&self, src: &SrcOperand) -> i32 {
346        match src {
347            SrcOperand::Reg(r) => self.get_register(*r),
348            SrcOperand::Uniform(slot) => {
349                let s = *slot as usize;
350                if s < self.uniform_values.len() {
351                    match &self.uniform_values[s] {
352                        UniformValue::Int(v) => *v,
353                        UniformValue::Vec2(x, _) => *x,
354                        UniformValue::Vec4(x, _, _, _) => *x,
355                        UniformValue::Mat4(m) => m[0],
356                    }
357                } else {
358                    0
359                }
360            }
361            SrcOperand::Immediate(v) => *v,
362        }
363    }
364
365    /// Execute a shader program.
366    ///
367    /// Registers should be pre-loaded with input values.
368    /// After execution, output registers contain the results.
369    pub fn execute(&mut self, program: &ShaderProgram) {
370        for instr in &program.instructions {
371            match instr {
372                TgsiInstruction::MOV { dst, src } => {
373                    let v = self.read_src(src);
374                    self.registers[*dst as usize % MAX_REGISTERS] = v;
375                }
376                TgsiInstruction::ADD { dst, src0, src1 } => {
377                    let a = self.read_src(src0);
378                    let b = self.read_src(src1);
379                    self.registers[*dst as usize % MAX_REGISTERS] = a.saturating_add(b);
380                }
381                TgsiInstruction::MUL { dst, src0, src1 } => {
382                    let a = self.read_src(src0);
383                    let b = self.read_src(src1);
384                    self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
385                }
386                TgsiInstruction::MAD {
387                    dst,
388                    src0,
389                    src1,
390                    src2,
391                } => {
392                    let a = self.read_src(src0);
393                    let b = self.read_src(src1);
394                    let c = self.read_src(src2);
395                    self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b).saturating_add(c);
396                }
397                TgsiInstruction::DP3 { dst, src0, src1 } => {
398                    // Dot product of first 3 components (uses consecutive registers)
399                    let a = self.read_src(src0);
400                    let b = self.read_src(src1);
401                    self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
402                }
403                TgsiInstruction::DP4 { dst, src0, src1 } => {
404                    let a = self.read_src(src0);
405                    let b = self.read_src(src1);
406                    self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
407                }
408                TgsiInstruction::TEX { dst, coord } => {
409                    // Stub: return the coordinate as a grey value
410                    let c = self.read_src(coord);
411                    self.registers[*dst as usize % MAX_REGISTERS] = c;
412                }
413                TgsiInstruction::SAMPLE {
414                    dst,
415                    coord,
416                    sampler: _,
417                } => {
418                    let c = self.read_src(coord);
419                    self.registers[*dst as usize % MAX_REGISTERS] = c;
420                }
421                TgsiInstruction::RCP { dst, src } => {
422                    let v = self.read_src(src);
423                    self.registers[*dst as usize % MAX_REGISTERS] = fp_rcp(v);
424                }
425                TgsiInstruction::RSQ { dst, src } => {
426                    let v = self.read_src(src);
427                    // Integer approximation of 1/sqrt(x)
428                    // Using Newton's method with one iteration
429                    if v <= 0 {
430                        self.registers[*dst as usize % MAX_REGISTERS] = 0;
431                    } else {
432                        // Rough initial guess
433                        let mut guess = FP_ONE;
434                        let mut test = fp_to_int(v);
435                        while test > 1 {
436                            test >>= 2;
437                            guess >>= 1;
438                        }
439                        if guess == 0 {
440                            guess = 1;
441                        }
442                        self.registers[*dst as usize % MAX_REGISTERS] = fp_rcp(guess);
443                    }
444                }
445            }
446        }
447    }
448
449    /// Execute a fragment shader for a single pixel, returning ARGB8888.
450    ///
451    /// Input: registers 0-3 are R, G, B, A in 16.16 fixed-point (0..255 range).
452    pub fn execute_fragment(&mut self, program: &ShaderProgram) -> u32 {
453        self.execute(program);
454
455        let r = fp_to_int(self.registers[0]).clamp(0, 255) as u32;
456        let g = fp_to_int(self.registers[1]).clamp(0, 255) as u32;
457        let b = fp_to_int(self.registers[2]).clamp(0, 255) as u32;
458        let a = fp_to_int(self.registers[3]).clamp(0, 255) as u32;
459
460        (a << 24) | (r << 16) | (g << 8) | b
461    }
462}
463
464// ---------------------------------------------------------------------------
465// Tests
466// ---------------------------------------------------------------------------
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471
472    #[test]
473    fn test_fp_mul() {
474        let a = fp_from_int(3);
475        let b = fp_from_int(4);
476        assert_eq!(fp_to_int(fp_mul(a, b)), 12);
477    }
478
479    #[test]
480    fn test_fp_rcp() {
481        let v = fp_from_int(2);
482        let r = fp_rcp(v);
483        // Should be approximately 0.5 in 16.16 = 32768
484        assert!((r - (FP_ONE / 2)).abs() < 2);
485    }
486
487    #[test]
488    fn test_shader_set_color() {
489        let prog = ShaderCompiler::compile_fragment(
490            "test",
491            &[ShaderOp::SetColor {
492                r: 128,
493                g: 64,
494                b: 32,
495                a: 255,
496            }],
497        );
498        assert!(prog.instruction_count() >= 4);
499        let mut exec = ShaderExecutor::new();
500        let pixel = exec.execute_fragment(&prog);
501        let r = (pixel >> 16) & 0xFF;
502        let g = (pixel >> 8) & 0xFF;
503        let b = pixel & 0xFF;
504        assert_eq!(r, 128);
505        assert_eq!(g, 64);
506        assert_eq!(b, 32);
507    }
508
509    #[test]
510    fn test_shader_passthrough() {
511        let prog = ShaderCompiler::compile_fragment("pass", &[ShaderOp::Passthrough]);
512        let mut exec = ShaderExecutor::new();
513        exec.set_register(0, fp_from_int(200));
514        exec.set_register(1, fp_from_int(100));
515        exec.set_register(2, fp_from_int(50));
516        exec.set_register(3, fp_from_int(255));
517        let pixel = exec.execute_fragment(&prog);
518        let r = (pixel >> 16) & 0xFF;
519        assert_eq!(r, 200);
520    }
521
522    #[test]
523    fn test_vertex_passthrough() {
524        let prog = ShaderCompiler::compile_vertex_passthrough("vtx");
525        assert_eq!(prog.shader_type, ShaderType::Vertex);
526        assert_eq!(prog.instruction_count(), 1);
527    }
528
529    #[test]
530    fn test_mov_instruction() {
531        let mut exec = ShaderExecutor::new();
532        let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
533        prog.push(TgsiInstruction::MOV {
534            dst: 0,
535            src: SrcOperand::Immediate(fp_from_int(42)),
536        });
537        exec.execute(&prog);
538        assert_eq!(fp_to_int(exec.get_register(0)), 42);
539    }
540
541    #[test]
542    fn test_add_instruction() {
543        let mut exec = ShaderExecutor::new();
544        exec.set_register(0, fp_from_int(10));
545        let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
546        prog.push(TgsiInstruction::ADD {
547            dst: 1,
548            src0: SrcOperand::Reg(0),
549            src1: SrcOperand::Immediate(fp_from_int(5)),
550        });
551        exec.execute(&prog);
552        assert_eq!(fp_to_int(exec.get_register(1)), 15);
553    }
554
555    #[test]
556    fn test_uniform_binding() {
557        let mut exec = ShaderExecutor::new();
558        exec.set_uniform(0, UniformValue::Int(fp_from_int(2)));
559        exec.set_register(0, fp_from_int(100));
560        let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
561        prog.bind_uniform("scale", 0);
562        prog.push(TgsiInstruction::MUL {
563            dst: 0,
564            src0: SrcOperand::Reg(0),
565            src1: SrcOperand::Uniform(0),
566        });
567        exec.execute(&prog);
568        assert_eq!(fp_to_int(exec.get_register(0)), 200);
569    }
570}