naga/proc/overloads/
mathfunction.rs

1//! Overload sets for [`ir::MathFunction`].
2
3use crate::proc::overloads::any_overload_set::AnyOverloadSet;
4use crate::proc::overloads::list::List;
5use crate::proc::overloads::regular::regular;
6use crate::proc::overloads::utils::{
7    float_scalars, float_scalars_unimplemented_abstract, list, pairs, rule, scalar_or_vecn, triples,
8};
9use crate::proc::overloads::OverloadSet;
10use crate::proc::type_methods::{concrete_int_scalars, vector_sizes};
11
12use crate::ir;
13
14impl ir::MathFunction {
15    pub fn overloads(self) -> impl OverloadSet {
16        use ir::MathFunction as Mf;
17
18        let set: AnyOverloadSet = match self {
19            // Component-wise unary numeric operations
20            Mf::Abs => regular!(1, SCALAR|VECN of NUMERIC).into(),
21            Mf::Sign => regular!(1, SCALAR|VECN of SIGNED_NUMERIC).into(),
22
23            // Component-wise binary numeric operations
24            Mf::Min | Mf::Max => regular!(2, SCALAR|VECN of NUMERIC).into(),
25
26            // Component-wise ternary numeric operations
27            Mf::Clamp => regular!(3, SCALAR|VECN of NUMERIC).into(),
28
29            // Component-wise unary floating-point operations
30            Mf::Sin
31            | Mf::Cos
32            | Mf::Tan
33            | Mf::Asin
34            | Mf::Acos
35            | Mf::Atan
36            | Mf::Sinh
37            | Mf::Cosh
38            | Mf::Tanh
39            | Mf::Asinh
40            | Mf::Acosh
41            | Mf::Atanh
42            | Mf::Saturate
43            | Mf::Radians
44            | Mf::Degrees
45            | Mf::Ceil
46            | Mf::Floor
47            | Mf::Round
48            | Mf::Fract
49            | Mf::Trunc
50            | Mf::Exp
51            | Mf::Exp2
52            | Mf::Log
53            | Mf::Log2
54            | Mf::Sqrt
55            | Mf::InverseSqrt => regular!(1, SCALAR|VECN of FLOAT).into(),
56
57            // Component-wise binary floating-point operations
58            Mf::Atan2 | Mf::Pow | Mf::Step => regular!(2, SCALAR|VECN of FLOAT).into(),
59
60            // Component-wise ternary floating-point operations
61            Mf::Fma | Mf::SmoothStep => regular!(3, SCALAR|VECN of FLOAT).into(),
62
63            // Component-wise unary concrete integer operations
64            Mf::CountTrailingZeros
65            | Mf::CountLeadingZeros
66            | Mf::CountOneBits
67            | Mf::ReverseBits
68            | Mf::FirstTrailingBit
69            | Mf::FirstLeadingBit => regular!(1, SCALAR|VECN of CONCRETE_INTEGER).into(),
70
71            // Packing functions
72            Mf::Pack4x8snorm | Mf::Pack4x8unorm => regular!(1, VEC4 of F32 -> U32).into(),
73            Mf::Pack2x16snorm | Mf::Pack2x16unorm | Mf::Pack2x16float => {
74                regular!(1, VEC2 of F32 -> U32).into()
75            }
76            Mf::Pack4xI8 => regular!(1, VEC4 of I32 -> U32).into(),
77            Mf::Pack4xU8 => regular!(1, VEC4 of U32 -> U32).into(),
78            Mf::Pack4xI8Clamp => regular!(1, VEC4 of I32 -> U32).into(),
79            Mf::Pack4xU8Clamp => regular!(1, VEC4 of U32 -> U32).into(),
80
81            // Unpacking functions
82            Mf::Unpack4x8snorm | Mf::Unpack4x8unorm => regular!(1, SCALAR of U32 -> Vec4F).into(),
83            Mf::Unpack2x16snorm | Mf::Unpack2x16unorm | Mf::Unpack2x16float => {
84                regular!(1, SCALAR of U32 -> Vec2F).into()
85            }
86            Mf::Unpack4xI8 => regular!(1, SCALAR of U32 -> Vec4I).into(),
87            Mf::Unpack4xU8 => regular!(1, SCALAR of U32 -> Vec4U).into(),
88            Mf::Dot4I8Packed => regular!(2, SCALAR of U32 -> I32).into(),
89            Mf::Dot4U8Packed => regular!(2, SCALAR of U32 -> U32).into(),
90
91            // One-off operations
92            Mf::Dot => regular!(2, VECN of NUMERIC -> Scalar).into(),
93            Mf::Modf => regular!(1, SCALAR|VECN of FLOAT_ABSTRACT_UNIMPLEMENTED -> Modf).into(),
94            Mf::Frexp => regular!(1, SCALAR|VECN of FLOAT_ABSTRACT_UNIMPLEMENTED -> Frexp).into(),
95            Mf::Ldexp => ldexp().into(),
96            Mf::Outer => outer().into(),
97            Mf::Cross => regular!(2, VEC3 of FLOAT).into(),
98            Mf::Distance => {
99                regular!(2, SCALAR|VECN of FLOAT_ABSTRACT_UNIMPLEMENTED -> Scalar).into()
100            }
101            Mf::Length => regular!(1, SCALAR|VECN of FLOAT -> Scalar).into(),
102            Mf::Normalize => regular!(1, VECN of FLOAT).into(),
103            Mf::FaceForward => regular!(3, VECN of FLOAT_ABSTRACT_UNIMPLEMENTED).into(),
104            Mf::Reflect => regular!(2, VECN of FLOAT_ABSTRACT_UNIMPLEMENTED).into(),
105            Mf::Refract => refract().into(),
106            Mf::Mix => mix().into(),
107            Mf::Inverse => regular!(1, MAT2X2|MAT3X3|MAT4X4 of FLOAT).into(),
108            Mf::Transpose => transpose().into(),
109            Mf::Determinant => regular!(1, MAT2X2|MAT3X3|MAT4X4 of FLOAT -> Scalar).into(),
110            Mf::QuantizeToF16 => regular!(1, SCALAR|VECN of F32).into(),
111            Mf::ExtractBits => extract_bits().into(),
112            Mf::InsertBits => insert_bits().into(),
113        };
114
115        set
116    }
117}
118
119fn ldexp() -> List {
120    /// Construct the exponent scalar given the mantissa's inner.
121    fn exponent_from_mantissa(mantissa: ir::Scalar) -> ir::Scalar {
122        match mantissa.kind {
123            ir::ScalarKind::AbstractFloat => ir::Scalar::ABSTRACT_INT,
124            ir::ScalarKind::Float => ir::Scalar::I32,
125            _ => unreachable!("not a float scalar"),
126        }
127    }
128
129    list(
130        // The ldexp mantissa argument can be any floating-point type.
131        float_scalars_unimplemented_abstract().flat_map(|mantissa_scalar| {
132            // The exponent type is the integer counterpart of the mantissa type.
133            let exponent_scalar = exponent_from_mantissa(mantissa_scalar);
134            // There are scalar and vector component-wise overloads.
135            scalar_or_vecn(mantissa_scalar)
136                .zip(scalar_or_vecn(exponent_scalar))
137                .map(move |(mantissa, exponent)| {
138                    let result = mantissa.clone();
139                    rule([mantissa, exponent], result)
140                })
141        }),
142    )
143}
144
145fn outer() -> List {
146    list(
147        triples(
148            vector_sizes(),
149            vector_sizes(),
150            float_scalars_unimplemented_abstract(),
151        )
152        .map(|(cols, rows, scalar)| {
153            let left = ir::TypeInner::Vector { size: cols, scalar };
154            let right = ir::TypeInner::Vector { size: rows, scalar };
155            let result = ir::TypeInner::Matrix {
156                columns: cols,
157                rows,
158                scalar,
159            };
160            rule([left, right], result)
161        }),
162    )
163}
164
165fn refract() -> List {
166    list(
167        pairs(vector_sizes(), float_scalars_unimplemented_abstract()).map(|(size, scalar)| {
168            let incident = ir::TypeInner::Vector { size, scalar };
169            let normal = incident.clone();
170            let ratio = ir::TypeInner::Scalar(scalar);
171            let result = incident.clone();
172            rule([incident, normal, ratio], result)
173        }),
174    )
175}
176
177fn transpose() -> List {
178    list(
179        triples(vector_sizes(), vector_sizes(), float_scalars()).map(|(a, b, scalar)| {
180            let input = ir::TypeInner::Matrix {
181                columns: a,
182                rows: b,
183                scalar,
184            };
185            let output = ir::TypeInner::Matrix {
186                columns: b,
187                rows: a,
188                scalar,
189            };
190            rule([input], output)
191        }),
192    )
193}
194
195fn extract_bits() -> List {
196    list(concrete_int_scalars().flat_map(|scalar| {
197        scalar_or_vecn(scalar).map(|input| {
198            let offset = ir::TypeInner::Scalar(ir::Scalar::U32);
199            let count = ir::TypeInner::Scalar(ir::Scalar::U32);
200            let output = input.clone();
201            rule([input, offset, count], output)
202        })
203    }))
204}
205
206fn insert_bits() -> List {
207    list(concrete_int_scalars().flat_map(|scalar| {
208        scalar_or_vecn(scalar).map(|input| {
209            let newbits = input.clone();
210            let offset = ir::TypeInner::Scalar(ir::Scalar::U32);
211            let count = ir::TypeInner::Scalar(ir::Scalar::U32);
212            let output = input.clone();
213            rule([input, newbits, offset, count], output)
214        })
215    }))
216}
217
218fn mix() -> List {
219    list(float_scalars().flat_map(|scalar| {
220        scalar_or_vecn(scalar).flat_map(move |input| {
221            let scalar_ratio = ir::TypeInner::Scalar(scalar);
222            [
223                rule([input.clone(), input.clone(), input.clone()], input.clone()),
224                rule([input.clone(), input.clone(), scalar_ratio], input),
225            ]
226        })
227    }))
228}
naga/proc/overloads/mathfunction.rs

naga/proc/overloads/
mathfunction.rs
