tfhe_hpu_backend/interface/
cmd.rs

1///
2/// Help with IOp management over HPU
3/// Track IOp status and handle backward update of associated HpuVariable
4use super::*;
5use crate::asm::iop::{Immediat, Operand, OperandKind};
6use crate::asm::{IOp, IOpcode};
7use variable::HpuVarWrapped;
8
9/// Underlying type used for Immediat value;
10pub type HpuImm = u128;
11
12/// Structure that hold an IOp with there associated operands
13/// Wrap operands memory with the IOp for proper lifetime management
14pub struct HpuCmd {
15    pub(crate) op: IOp,
16    pub(crate) dst: Vec<HpuVarWrapped>,
17    pub(crate) src: Vec<HpuVarWrapped>,
18    // NB: No need to track Immediat lifetime. It's simply constant completely held by the IOp
19    // definition
20}
21
22impl HpuCmd {
23    pub fn new(
24        opcode: IOpcode,
25        dst: &[HpuVarWrapped],
26        src: &[HpuVarWrapped],
27        imm: &[HpuImm],
28    ) -> Self {
29        // TODO Check that dst/rhs_x backend match
30        // Check arguments compliance with IOp prototype if any
31        #[cfg(debug_assertions)]
32        if let Some(format) = crate::asm::iop::IOP_LUT.hex.get(&opcode) {
33            assert_eq!(
34                dst.len(),
35                format.proto.dst.len(),
36                "Error {}: Invalid number of dst arguments",
37                format.name
38            );
39            assert_eq!(
40                src.len(),
41                format.proto.src.len(),
42                "Error {}: Invalid number of dst arguments",
43                format.name
44            );
45            assert_eq!(
46                imm.len(),
47                format.proto.imm,
48                "Error {}: Invalid number of dst arguments",
49                format.name
50            );
51        }
52
53        // Extract Operands definition from HpuVar
54        let dst_op = dst
55            .iter()
56            .map(|var| {
57                Operand::new(
58                    var.width as u8,
59                    var.id.0 as u16,
60                    1, /* TODO handle vec source !? */
61                    Some(OperandKind::Dst),
62                )
63            })
64            .collect::<Vec<_>>();
65        let src_op = src
66            .iter()
67            .map(|var| {
68                Operand::new(
69                    var.width as u8,
70                    var.id.0 as u16,
71                    1, /* TODO handle vec source !? */
72                    Some(OperandKind::Src),
73                )
74            })
75            .collect::<Vec<_>>();
76        let imm_op = imm
77            .iter()
78            .map(|var| Immediat::from_cst(*var))
79            .collect::<Vec<_>>();
80
81        let op = IOp::new(opcode, dst_op, src_op, imm_op);
82        // TODO set op_width
83
84        let dst = dst
85            .iter()
86            .map(|var| {
87                // Update dst state to OpPending
88                var.inner.lock().unwrap().operation_pending();
89                (*var).clone()
90            })
91            .collect::<Vec<_>>();
92        let src = src.iter().map(|var| (*var).clone()).collect::<Vec<_>>();
93        Self { op, dst, src }
94    }
95
96    pub fn op(&self) -> &IOp {
97        &self.op
98    }
99}
100
101/// Generic interface
102impl HpuCmd {
103    pub fn exec_raw(
104        opcode: crate::asm::IOpcode,
105        dst: &[HpuVarWrapped],
106        rhs_ct: &[HpuVarWrapped],
107        rhs_imm: &[HpuImm],
108    ) {
109        // Create associated command
110        let cmd = Self::new(opcode, dst, rhs_ct, rhs_imm);
111        // Issue it on Hpubackend
112        dst.first()
113            .expect("Try to generate an IOp without any destination")
114            .cmd_api
115            .send(cmd)
116            .expect("Issue with cmd_api");
117    }
118
119    // TODO add more runtime check on prototype ?
120    pub fn exec(
121        proto: &crate::asm::iop::IOpProto,
122        opcode: crate::asm::IOpcode,
123        rhs_ct: &[HpuVarWrapped],
124        rhs_imm: &[HpuImm],
125    ) -> Vec<HpuVarWrapped> {
126        let dst = proto
127            .dst
128            .iter()
129            .map(|m| rhs_ct[0].fork(*m))
130            .collect::<Vec<_>>();
131        Self::exec_raw(opcode, &dst, rhs_ct, rhs_imm);
132        dst
133    }
134
135    pub fn exec_assign(
136        proto: &crate::asm::iop::IOpProto,
137        opcode: crate::asm::IOpcode,
138        rhs_ct: &[HpuVarWrapped],
139        rhs_imm: &[HpuImm],
140    ) {
141        // Clone dst sub-array from srcs
142        let dst = std::iter::zip(proto.dst.iter(), rhs_ct.iter())
143            .map(|(p, v)| {
144                debug_assert_eq!(
145                    *p, v.mode,
146                    "Assign with invalid prototype, rhs mode don't match"
147                );
148                v.clone()
149            })
150            .collect::<Vec<_>>();
151        Self::exec_raw(opcode, &dst, rhs_ct, rhs_imm);
152    }
153}