zhc_ir 0.8.0

Graph-based intermediate representation framework with dialect support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use std::{hash::Hash, rc::Rc};

use zhc_utils::{small::SmallVec, svec};

use crate::{
    Dialect, DialectInstructionSet, Format, FormatContext, IR, OpId, OpRef, Signature, ValId, sig,
    visualization::{
        layoutlang::{LayoutDialect, LayoutTypeSystem},
        visual_annotation::VisualAnnotation,
    },
};

#[derive(Debug)]
pub struct OpContent {
    pub args: SmallVec<String>,
    pub returns: SmallVec<String>,
    pub call: String,
    pub comment: Option<String>,
    pub annotation: Option<Rc<dyn VisualAnnotation>>,
}

impl Clone for OpContent {
    fn clone(&self) -> Self {
        Self {
            args: self.args.clone(),
            returns: self.returns.clone(),
            call: self.call.clone(),
            comment: self.comment.clone(),
            annotation: self.annotation.clone(),
        }
    }
}

impl PartialEq for OpContent {
    fn eq(&self, other: &Self) -> bool {
        // Annotation is intentionally excluded from equality
        self.args == other.args
            && self.returns == other.returns
            && self.call == other.call
            && self.comment == other.comment
    }
}

impl Eq for OpContent {}

impl Hash for OpContent {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        // Annotation is intentionally excluded from hashing
        self.args.hash(state);
        self.returns.hash(state);
        self.call.hash(state);
        self.comment.hash(state);
    }
}

impl OpContent {
    pub fn from_op<'ir, D: Dialect>(opref: &OpRef<'ir, D>, ctx: &FormatContext) -> Self {
        OpContent {
            args: opref
                .get_args_iter()
                .map(|a| a.fmt_to_string(ctx))
                .collect(),
            returns: opref
                .get_returns_iter()
                .map(|a| a.fmt_to_string(ctx))
                .collect(),
            call: opref.fmt_to_string(&ctx.clone().show_comments(false).show_types(false)),
            comment: None,
            annotation: None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LayoutInstructionSet {
    Operation {
        opid: OpId,
        op: OpContent,
        args: SmallVec<ValId>,
        returns: SmallVec<ValId>,
    },
    Dummy {
        valid: ValId,
    },
    Group {
        ir: IR<LayoutDialect>,
        name: String,
    },
    GroupInput {
        pos: u16,
        valid: ValId,
    },
    GroupOutput {
        pos: u16,
        valid: ValId,
    },
}

impl Format for LayoutInstructionSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, ctx: &crate::FormatContext) -> std::fmt::Result {
        match self {
            LayoutInstructionSet::Operation { opid, .. } => write!(f, "operation<{opid}>"),
            LayoutInstructionSet::Dummy { valid, .. } => write!(f, "dummy<{valid}>"),
            LayoutInstructionSet::Group { ir, name, .. } => {
                let inner_ctx = ctx.with_prefix("    ").with_next_nested_prefix();
                writeln!(f, "group<\"{}\"> {{", name)?;
                Format::fmt(ir, f, &inner_ctx)?;
                write!(f, "\n{}}}", ctx.prefix())
            }
            LayoutInstructionSet::GroupInput { pos, .. } => write!(f, "group_input<{pos}>"),
            LayoutInstructionSet::GroupOutput { pos, .. } => write!(f, "group_output<{pos}>"),
        }
    }
}

impl DialectInstructionSet for LayoutInstructionSet {
    type TypeSystem = LayoutTypeSystem;

    fn get_signature(&self) -> crate::Signature<Self::TypeSystem> {
        match self {
            LayoutInstructionSet::Operation { op, .. } => Signature(
                svec![LayoutTypeSystem::Value; op.args.len()],
                svec![LayoutTypeSystem::Value; op.returns.len()],
            ),
            LayoutInstructionSet::Dummy { .. } => {
                sig![(LayoutTypeSystem::Value) -> (LayoutTypeSystem::Value)]
            }
            LayoutInstructionSet::Group { ir, .. } => {
                let n_inputs = ir
                    .walk_ops_linear()
                    .filter(|op| {
                        matches!(
                            op.get_instruction(),
                            LayoutInstructionSet::GroupInput { .. }
                        )
                    })
                    .count();
                let n_outputs = ir
                    .walk_ops_linear()
                    .filter(|op| {
                        matches!(
                            op.get_instruction(),
                            LayoutInstructionSet::GroupOutput { .. }
                        )
                    })
                    .count();
                Signature(
                    svec![LayoutTypeSystem::Value; n_inputs],
                    svec![LayoutTypeSystem::Value; n_outputs],
                )
            }
            LayoutInstructionSet::GroupInput { .. } => sig![() -> (LayoutTypeSystem::Value)],
            LayoutInstructionSet::GroupOutput { .. } => sig![(LayoutTypeSystem::Value) -> ()],
        }
    }
}