Skip to main content

gaia_assembler/backends/
gcn.rs

1use crate::{
2    backends::{Backend, GeneratedFiles},
3    config::GaiaConfig,
4    instruction::{DomainInstruction, GaiaInstruction},
5    program::{GaiaModule, GaiaTerminator},
6};
7use gaia_types::{
8    helpers::{AbiCompatible, ApiCompatible, Architecture, CompilationTarget},
9    neural::NeuralNode,
10    Result,
11};
12use gcn_assembler::{
13    instructions::{GcnInstruction, GcnReg},
14    program::{GcnKernel, GcnProgram},
15    GcnWriter,
16};
17use std::collections::HashMap;
18
19pub struct GcnBackend {
20    writer: GcnWriter,
21}
22
23impl GcnBackend {
24    pub fn new() -> Self {
25        Self { writer: GcnWriter::new() }
26    }
27}
28
29impl Backend for GcnBackend {
30    fn name(&self) -> &'static str {
31        "AMD GCN"
32    }
33
34    fn primary_target(&self) -> CompilationTarget {
35        CompilationTarget {
36            build: Architecture::RISCV64, // Placeholder
37            host: AbiCompatible::GCN,
38            target: ApiCompatible::Unknown,
39        }
40    }
41
42    fn match_score(&self, target: &CompilationTarget) -> f32 {
43        if target.build == Architecture::RISCV64 {
44            return 50.0;
45        }
46        0.0
47    }
48
49    fn generate(&self, program: &GaiaModule, _config: &GaiaConfig) -> Result<GeneratedFiles> {
50        let mut files = HashMap::new();
51
52        // 转换 GaiaModule -> GcnProgram
53        let gcn_program = convert_gaia_to_gcn(program)?;
54
55        // 写入二进制
56        let binary = self.writer.write(&gcn_program)?;
57        files.insert(format!("{}.hsaco", program.name), binary);
58
59        Ok(GeneratedFiles { files, diagnostics: vec![] })
60    }
61}
62
63fn convert_gaia_to_gcn(module: &GaiaModule) -> Result<GcnProgram> {
64    let mut gcn_program = GcnProgram::new(module.name.clone());
65
66    for function in &module.functions {
67        let mut instructions = Vec::new();
68        for block in &function.blocks {
69            for inst in &block.instructions {
70                match inst {
71                    GaiaInstruction::Domain(DomainInstruction::Neural(node)) => {
72                        // 映射神经网络算子到 GCN 指令
73                        if let NeuralNode::Convolution { .. } = node {
74                            instructions.push(GcnInstruction::VDot2F32F16 {
75                                dst: GcnReg::VGPR(0),
76                                src0: GcnReg::VGPR(1),
77                                src1: GcnReg::VGPR(2),
78                            });
79                        }
80                    }
81                    _ => {
82                        instructions.push(GcnInstruction::SNop(0));
83                    }
84                }
85            }
86
87            match &block.terminator {
88                GaiaTerminator::Return => {
89                    instructions.push(GcnInstruction::SEndPgm);
90                }
91                _ => {}
92            }
93        }
94
95        gcn_program.kernels.push(GcnKernel { name: function.name.clone(), instructions, args: Vec::new() });
96    }
97
98    Ok(gcn_program)
99}