tfhe_hpu_backend/asm/
mod.rs

1pub mod dop;
2pub use dop::arg::Arg as DOpArg;
3pub use dop::{DOp, DigitParameters, ImmId, MemId, Pbs, PbsGid, PbsLut, RegId, ToHex};
4pub mod iop;
5pub use iop::{AsmIOpcode, IOp, IOpProto, IOpcode, OperandKind};
6
7use std::collections::VecDeque;
8use std::io::{BufRead, Write};
9
10pub const ASM_COMMENT_PREFIX: [char; 2] = [';', '#'];
11
12// Common type used in both DOp/IOp definition --------------------------------
13/// Ciphertext Id
14/// On-board memory is viewed as an array of ciphertext,
15/// Thus, instead of using bytes address, ct id is used
16/// => Id of the first ciphertext of the vector
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
18pub struct CtId(pub u16);
19
20// ---------------------------------------------------------------------------
21
22/// Simple test for Asm parsing
23#[cfg(test)]
24mod tests;
25
26/// Type to aggregate Op and header
27/// Aim is to kept correct interleaving while parsing
28#[derive(Debug, Clone)]
29pub enum AsmOp<Op> {
30    Comment(String),
31    Stmt(Op),
32}
33
34impl<Op: dop::arg::ToFlush> AsmOp<Op> {
35    pub fn to_flush(&mut self) {
36        if let AsmOp::Stmt(op) = self {
37            *op = op.to_flush();
38        }
39    }
40}
41
42impl<Op: std::fmt::Display> std::fmt::Display for AsmOp<Op> {
43    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
44        match self {
45            Self::Comment(c) => write!(f, "{}{c}", ASM_COMMENT_PREFIX[0]),
46            Self::Stmt(op) => write!(f, "{op}"),
47        }
48    }
49}
50
51/// Generic struct to represent sequence of operations
52/// Used to extract OP from ASM file
53/// Work on any kind of Op that implement FromStr
54#[derive(Debug, Clone)]
55pub struct Program<Op>(Vec<AsmOp<Op>>);
56
57impl<Op> Default for Program<Op> {
58    fn default() -> Self {
59        Self(Vec::new())
60    }
61}
62
63impl<Op> Program<Op> {
64    pub fn new(ops: Vec<AsmOp<Op>>) -> Self {
65        Self(ops)
66    }
67    /// Push a new statement in the program
68    pub fn push_stmt(&mut self, op: Op) {
69        self.0.push(AsmOp::Stmt(op))
70    }
71    /// Push a new statement in the program
72    /// Returns the position in which the statement was inserted
73    pub fn push_stmt_pos(&mut self, op: Op) -> usize {
74        let ret = self.0.len();
75        self.0.push(AsmOp::Stmt(op));
76        ret
77    }
78    /// Push a new comment in the program
79    pub fn push_comment(&mut self, comment: String) {
80        self.0.push(AsmOp::Comment(comment))
81    }
82
83    pub fn get_stmt_mut(&mut self, i: usize) -> &mut AsmOp<Op> {
84        &mut self.0[i]
85    }
86}
87
88impl<Op> std::ops::Deref for Program<Op> {
89    type Target = Vec<AsmOp<Op>>;
90
91    fn deref(&self) -> &Self::Target {
92        &self.0
93    }
94}
95
96impl<Op: std::fmt::Display> std::fmt::Display for Program<Op> {
97    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
98        for op in self.0.iter() {
99            writeln!(f, "{op}")?;
100        }
101        Ok(())
102    }
103}
104
105impl<Op, Err> Program<Op>
106where
107    Op: std::str::FromStr<Err = Err>,
108    Err: std::error::Error,
109{
110    /// Generic function to extract OP from ASM file
111    /// Work on any kind of Op that implement FromStr
112    pub fn read_asm(file: &str) -> Result<Self, anyhow::Error> {
113        // Open file
114        let rd_f = std::io::BufReader::new(
115            std::fs::OpenOptions::new()
116                .create(false)
117                .read(true)
118                .open(file)?,
119        );
120
121        let mut asm_ops = Vec::new();
122        for (line, val) in rd_f.lines().map_while(Result::ok).enumerate() {
123            if let Some(comment) = val.trim().strip_prefix(ASM_COMMENT_PREFIX) {
124                asm_ops.push(AsmOp::Comment(comment.to_string()))
125            } else if !val.is_empty() {
126                match Op::from_str(&val) {
127                    Ok(op) => asm_ops.push(AsmOp::Stmt(op)),
128                    Err(err) => {
129                        tracing::warn!("ReadAsm failed @{file}:{}", line + 1);
130                        anyhow::bail!("ReadAsm failed @{file}:{} with {}", line + 1, err);
131                    }
132                }
133            }
134        }
135        Ok(Self(asm_ops))
136    }
137}
138
139impl<Op> Program<Op>
140where
141    Op: std::fmt::Display,
142{
143    /// Generic function to write Op in ASM file
144    /// Work on any kind of Op that implement Display
145    pub fn write_asm(&self, file: &str) -> Result<(), anyhow::Error> {
146        // Create path
147        let path = std::path::Path::new(file);
148        if let Some(dir_p) = path.parent() {
149            std::fs::create_dir_all(dir_p).unwrap();
150        }
151
152        // Open file
153        let mut wr_f = std::fs::OpenOptions::new()
154            .create(true)
155            .write(true)
156            .truncate(true)
157            .open(path)?;
158
159        writeln!(wr_f, "{self}").map_err(anyhow::Error::new)
160    }
161}
162
163// Implement dedicated hex parser/dumper for DOp
164impl Program<dop::DOp> {
165    /// Generic function to extract OP from hex file
166    /// Work on any kind of Op that implement FromStr
167    pub fn read_hex(file: &str) -> Result<Self, anyhow::Error> {
168        // Open file
169        let rd_f = std::io::BufReader::new(
170            std::fs::OpenOptions::new()
171                .create(false)
172                .read(true)
173                .open(file)
174                .unwrap_or_else(|_| panic!("Invalid HEX file {file}")),
175        );
176
177        let mut prog = Self::default();
178        for (line, val) in rd_f.lines().map_while(Result::ok).enumerate() {
179            if let Some(comment) = val.trim().strip_prefix(ASM_COMMENT_PREFIX) {
180                prog.push_comment(comment.to_string());
181            } else {
182                let val_u32 =
183                    dop::DOpRepr::from_str_radix(std::str::from_utf8(val.as_bytes()).unwrap(), 16)?;
184                match dop::DOp::from_hex(val_u32) {
185                    Ok(op) => prog.push_stmt(op),
186                    Err(err) => {
187                        tracing::warn!("DOp::ReadHex failed @{file}:{}", line + 1);
188                        return Err(err.into());
189                    }
190                }
191            }
192        }
193        Ok(prog)
194    }
195
196    /// Generic function to write Op in Hex file
197    pub fn write_hex(&self, file: &str) -> Result<(), anyhow::Error> {
198        // Create path
199        let path = std::path::Path::new(file);
200        if let Some(dir_p) = path.parent() {
201            std::fs::create_dir_all(dir_p).unwrap();
202        }
203
204        // Open file
205        let mut wr_f = std::fs::OpenOptions::new()
206            .create(true)
207            .write(true)
208            .truncate(true)
209            .open(path)?;
210
211        for op in self.0.iter() {
212            match op {
213                AsmOp::Comment(comment) => writeln!(wr_f, "{}{}", ASM_COMMENT_PREFIX[0], comment)?,
214                AsmOp::Stmt(op) => writeln!(wr_f, "{:x}", op.to_hex())?,
215            }
216        }
217        Ok(())
218    }
219}
220
221impl Program<dop::DOp> {
222    /// Convert a program of Dops in translation table
223    pub fn tr_table(&self) -> Vec<dop::DOpRepr> {
224        let ops_stream = self
225            .iter()
226            .filter_map(|op| match op {
227                AsmOp::Comment(_) => None,
228                AsmOp::Stmt(op) => Some(op),
229            })
230            .collect::<Vec<_>>();
231
232        let mut words_stream = Vec::with_capacity(ops_stream.len() + 1);
233        // First word of the stream is length in DOp
234        words_stream.push(ops_stream.len() as u32);
235
236        ops_stream.iter().for_each(|op| {
237            words_stream.push(op.to_hex());
238        });
239        words_stream
240    }
241}
242
243// Implement dedicated hex parser/dumper for IOp
244impl Program<iop::IOp> {
245    /// Generic function to extract OP from hex file
246    pub fn read_hex(file: &str) -> Result<Self, anyhow::Error> {
247        // Open file
248        let rd_f = std::io::BufReader::new(
249            std::fs::OpenOptions::new()
250                .create(false)
251                .read(true)
252                .open(file)
253                .unwrap_or_else(|_| panic!("Invalid HEX file {file}")),
254        );
255
256        let mut prog = Self::default();
257        // Buffer word stream.
258        // When comment token occurred, convert the word stream into IOp
259        // -> No comment could be inserted in a middle of IOp word stream
260        let mut word_stream = VecDeque::new();
261        let mut file_len = 0;
262
263        for val in rd_f.lines().map_while(Result::ok) {
264            file_len += 1;
265            if let Some(comment) = val.trim().strip_prefix(ASM_COMMENT_PREFIX) {
266                while !word_stream.is_empty() {
267                    match iop::IOp::from_words(&mut word_stream) {
268                        Ok(op) => prog.push_stmt(op),
269                        Err(err) => {
270                            tracing::warn!(
271                                "IOp::ReadHex failed @{file}:{}",
272                                file_len - word_stream.len()
273                            );
274                            return Err(err.into());
275                        }
276                    }
277                }
278                prog.push_comment(comment.to_string());
279            } else {
280                let word = iop::IOpWordRepr::from_str_radix(
281                    std::str::from_utf8(val.as_bytes()).unwrap(),
282                    16,
283                )?;
284                word_stream.push_back(word);
285            }
286        }
287        // Flush word stream
288        while !word_stream.is_empty() {
289            match iop::IOp::from_words(&mut word_stream) {
290                Ok(op) => prog.push_stmt(op),
291                Err(err) => {
292                    tracing::warn!(
293                        "IOp::ReadHex failed @{file}:{}",
294                        file_len - word_stream.len()
295                    );
296                    return Err(err.into());
297                }
298            }
299        }
300        Ok(prog)
301    }
302
303    /// Generic function to write Op in Hex file
304    pub fn write_hex(&self, file: &str) -> Result<(), anyhow::Error> {
305        // Create path
306        let path = std::path::Path::new(file);
307        if let Some(dir_p) = path.parent() {
308            std::fs::create_dir_all(dir_p).unwrap();
309        }
310
311        // Open file
312        let mut wr_f = std::fs::OpenOptions::new()
313            .create(true)
314            .write(true)
315            .truncate(true)
316            .open(path)?;
317
318        for op in self.0.iter() {
319            match op {
320                AsmOp::Comment(comment) => writeln!(wr_f, "{}{}", ASM_COMMENT_PREFIX[0], comment)?,
321                AsmOp::Stmt(op) => {
322                    op.to_words()
323                        .into_iter()
324                        .try_for_each(|word| writeln!(wr_f, "{word:0>8x}"))?;
325                }
326            }
327        }
328        Ok(())
329    }
330}
tfhe_hpu_backend/asm/mod.rs

tfhe_hpu_backend/asm/
mod.rs
