pliron_math/ops/
definitions.rs

1//! Private operation definitions and their generated attribute storage.
2
3use crate::Diagnostic;
4use pliron::{
5    common_traits::Verify,
6    context::{Context, Ptr},
7    derive::pliron_op,
8    op::Op,
9    operation::Operation,
10    printable::Printable,
11    result::Result,
12    r#type::{TypeHandle, Typed},
13    value::Value,
14};
15use pliron_arith::{
16    attributes::{FastMathFlags, FastMathFlagsAttr},
17    shaped::{self, Shaped},
18};
19
20pliron::dict_key!(ATTR_KEY_MATH_FASTMATH, "math_fastmath");
21
22/// Verification error carrying the operation's name and source location.
23fn fail(ctx: &Context, op: Ptr<Operation>, reason: &'static str) -> pliron::result::Error {
24    let location = op.deref(ctx).loc().disp(ctx).to_string();
25    Diagnostic::InvalidOperation {
26        operation: Operation::get_op_dyn(op, ctx).get_opid().to_string(),
27        location,
28        reason,
29    }
30    .into()
31}
32
33/// Checks operand and result counts and rejects regions and successors.
34fn verify_structure(ctx: &Context, op: Ptr<Operation>, operands: usize) -> Result<()> {
35    let (num_operands, num_results, num_regions, num_successors) = {
36        let ir = op.deref(ctx);
37        (
38            ir.get_num_operands(),
39            ir.get_num_results(),
40            ir.num_regions(),
41            ir.get_num_successors(),
42        )
43    };
44    if num_operands != operands || num_results != 1 {
45        return Err(fail(ctx, op, "operand or result count mismatch"));
46    }
47    if num_regions != 0 || num_successors != 0 {
48        return Err(fail(
49            ctx,
50            op,
51            "operation must not have regions or successors",
52        ));
53    }
54    Ok(())
55}
56
57fn operand_type(ctx: &Context, op: Ptr<Operation>, index: usize) -> TypeHandle {
58    op.deref(ctx).get_operand(index).get_type(ctx)
59}
60
61fn result_type(ctx: &Context, op: Ptr<Operation>) -> TypeHandle {
62    op.deref(ctx).get_result(0).get_type(ctx)
63}
64
65fn classify(ctx: &Context, op: Ptr<Operation>, ty: TypeHandle) -> Result<Shaped> {
66    shaped::classify(ctx, ty).ok_or_else(|| {
67        fail(
68            ctx,
69            op,
70            "type must be a scalar or a tensor or vector of scalars",
71        )
72    })
73}
74
75/// Element-type family accepted by an operation.
76#[derive(Clone, Copy)]
77enum Domain {
78    /// Floating-point elements.
79    Float,
80    /// Signless integer elements; `math` integer operations reject `index`.
81    Integer,
82}
83
84fn in_domain(ctx: &Context, element: TypeHandle, domain: Domain) -> bool {
85    match domain {
86        Domain::Float => shaped::is_float(ctx, element),
87        Domain::Integer => shaped::signless_width(ctx, element).is_some(),
88    }
89}
90
91/// Verifies `SameOperandsAndResultType` over `arity` operands in `domain`.
92fn verify_same_type(ctx: &Context, op: Ptr<Operation>, arity: usize, domain: Domain) -> Result<()> {
93    verify_structure(ctx, op, arity)?;
94    let result = result_type(ctx, op);
95    if (0..arity).any(|i| operand_type(ctx, op, i) != result) {
96        return Err(fail(ctx, op, "operands and result must have the same type"));
97    }
98    let shaped = classify(ctx, op, result)?;
99    if !in_domain(ctx, shaped.element, domain) {
100        return Err(fail(
101            ctx,
102            op,
103            "element type is outside the operation's domain",
104        ));
105    }
106    Ok(())
107}
108
109fn new_op<T: Op>(
110    ctx: &mut Context,
111    make: fn(Ptr<Operation>) -> T,
112    result: TypeHandle,
113    operands: Vec<Value>,
114) -> T {
115    let info = T::get_concrete_op_info();
116    make(Operation::new(ctx, info, vec![result], operands, vec![], 0))
117}
118
119macro_rules! fastmath_accessors {
120    ($struct:ident) => {
121        pliron_arith::shared_attr!(
122            $struct,
123            ATTR_KEY_MATH_FASTMATH,
124            FastMathFlagsAttr,
125            get_attr_math_fastmath,
126            set_attr_math_fastmath
127        );
128        impl $struct {
129            /// Returns the fast-math flags, or [`FastMathFlags::NONE`] when absent.
130            pub fn fastmath(&self, ctx: &Context) -> FastMathFlags {
131                self.get_attr_math_fastmath(ctx)
132                    .map(|attr| attr.value)
133                    .unwrap_or_default()
134            }
135            /// Replaces the fast-math flags.
136            pub fn set_fastmath(&self, ctx: &Context, flags: FastMathFlags) {
137                self.set_attr_math_fastmath(ctx, FastMathFlagsAttr::new(flags));
138            }
139        }
140    };
141}
142
143macro_rules! result_accessor {
144    ($struct:ident) => {
145        impl $struct {
146            /// Returns the result.
147            pub fn result(&self, ctx: &Context) -> Value {
148                self.get_operation().deref(ctx).get_result(0)
149            }
150        }
151    };
152}
153
154macro_rules! unary_op {
155    ($struct:ident, $name:literal, $doc:literal, Float) => {
156        #[doc = $doc]
157        ///
158        /// Operand and result share one floating-point scalar, tensor, or vector type.
159        #[pliron_op(
160            name = $name,
161            format = "$0 attr_dict `:` typesig",
162            interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
163        )]
164        pub struct $struct;
165        unary_op!(@common $struct, Domain::Float);
166        fastmath_accessors!($struct);
167    };
168    ($struct:ident, $name:literal, $doc:literal, Integer) => {
169        #[doc = $doc]
170        ///
171        /// Operand and result share one signless-integer scalar, tensor, or vector type.
172        #[pliron_op(
173            name = $name,
174            format = "$0 attr_dict `:` typesig",
175            interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
176        )]
177        pub struct $struct;
178        unary_op!(@common $struct, Domain::Integer);
179    };
180    (@common $struct:ident, $domain:expr) => {
181        impl $struct {
182            /// Allocates the operation; the result type is the operand's type.
183            pub fn new(ctx: &mut Context, operand: Value) -> Self {
184                let ty = operand.get_type(ctx);
185                new_op(ctx, |op| Self { op }, ty, vec![operand])
186            }
187            /// Returns the operand.
188            pub fn operand(&self, ctx: &Context) -> Value {
189                self.get_operation().deref(ctx).get_operand(0)
190            }
191        }
192        result_accessor!($struct);
193        impl Verify for $struct {
194            fn verify(&self, ctx: &Context) -> Result<()> {
195                verify_same_type(ctx, self.get_operation(), 1, $domain)
196            }
197        }
198    };
199}
200
201macro_rules! binary_op {
202    ($struct:ident, $name:literal, $doc:literal, Float) => {
203        #[doc = $doc]
204        ///
205        /// Operands and result share one floating-point scalar, tensor, or vector type.
206        #[pliron_op(
207            name = $name,
208            format = "$0 `,` $1 attr_dict `:` typesig",
209            interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
210        )]
211        pub struct $struct;
212        binary_op!(@common $struct, Domain::Float);
213        fastmath_accessors!($struct);
214    };
215    ($struct:ident, $name:literal, $doc:literal, Integer) => {
216        #[doc = $doc]
217        ///
218        /// Operands and result share one signless-integer scalar, tensor, or vector type.
219        #[pliron_op(
220            name = $name,
221            format = "$0 `,` $1 attr_dict `:` typesig",
222            interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
223        )]
224        pub struct $struct;
225        binary_op!(@common $struct, Domain::Integer);
226    };
227    (@common $struct:ident, $domain:expr) => {
228        impl $struct {
229            /// Allocates the operation; the result type is the left operand's type.
230            pub fn new(ctx: &mut Context, lhs: Value, rhs: Value) -> Self {
231                let ty = lhs.get_type(ctx);
232                new_op(ctx, |op| Self { op }, ty, vec![lhs, rhs])
233            }
234            /// Returns the left operand.
235            pub fn lhs(&self, ctx: &Context) -> Value {
236                self.get_operation().deref(ctx).get_operand(0)
237            }
238            /// Returns the right operand.
239            pub fn rhs(&self, ctx: &Context) -> Value {
240                self.get_operation().deref(ctx).get_operand(1)
241            }
242        }
243        result_accessor!($struct);
244        impl Verify for $struct {
245            fn verify(&self, ctx: &Context) -> Result<()> {
246                verify_same_type(ctx, self.get_operation(), 2, $domain)
247            }
248        }
249    };
250}
251
252macro_rules! classification_op {
253    ($struct:ident, $name:literal, $doc:literal) => {
254        #[doc = $doc]
255        ///
256        /// The operand is floating-point; the result is `i1` with the operand's shape.
257        #[pliron_op(
258                            name = $name,
259                            format = "$0 attr_dict `:` typesig",
260                            interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
261                        )]
262        pub struct $struct;
263        impl $struct {
264            /// Allocates the operation with an `i1` result shaped like the operand.
265            ///
266            /// A non-shaped operand yields a scalar `i1`; call `verify` to reject it.
267            pub fn new(ctx: &mut Context, operand: Value) -> Self {
268                let ty = operand.get_type(ctx);
269                let i1 = shaped::i1(ctx);
270                let result = shaped::classify(ctx, ty)
271                    .map(|s| shaped::with_element(ctx, &s, i1))
272                    .unwrap_or(i1);
273                new_op(ctx, |op| Self { op }, result, vec![operand])
274            }
275            /// Returns the operand.
276            pub fn operand(&self, ctx: &Context) -> Value {
277                self.get_operation().deref(ctx).get_operand(0)
278            }
279        }
280        result_accessor!($struct);
281        fastmath_accessors!($struct);
282        impl Verify for $struct {
283            fn verify(&self, ctx: &Context) -> Result<()> {
284                let op = self.get_operation();
285                verify_structure(ctx, op, 1)?;
286                let input = classify(ctx, op, operand_type(ctx, op, 0))?;
287                if !shaped::is_float(ctx, input.element) {
288                    return Err(fail(ctx, op, "operand must be floating-point"));
289                }
290                let output = classify(ctx, op, result_type(ctx, op))?;
291                if !output.same_shape(&input) || output.element != shaped::i1(ctx) {
292                    return Err(fail(ctx, op, "result must be i1 with the operand's shape"));
293                }
294                Ok(())
295            }
296        }
297    };
298}
299
300unary_op!(AbsFOp, "math.absf", "Floating-point absolute value.", Float);
301unary_op!(
302    CeilOp,
303    "math.ceil",
304    "Rounds toward positive infinity.",
305    Float
306);
307unary_op!(
308    FloorOp,
309    "math.floor",
310    "Rounds toward negative infinity.",
311    Float
312);
313unary_op!(
314    RoundOp,
315    "math.round",
316    "Rounds to nearest, ties away from zero.",
317    Float
318);
319unary_op!(
320    RoundEvenOp,
321    "math.roundeven",
322    "Rounds to nearest, ties to even.",
323    Float
324);
325unary_op!(TruncOp, "math.trunc", "Rounds toward zero.", Float);
326unary_op!(ExpOp, "math.exp", "Base-e exponential.", Float);
327unary_op!(Exp2Op, "math.exp2", "Base-2 exponential.", Float);
328unary_op!(
329    ExpM1Op,
330    "math.expm1",
331    "`exp(x) - 1`, accurate near zero.",
332    Float
333);
334unary_op!(LogOp, "math.log", "Natural logarithm.", Float);
335unary_op!(Log2Op, "math.log2", "Base-2 logarithm.", Float);
336unary_op!(Log10Op, "math.log10", "Base-10 logarithm.", Float);
337unary_op!(
338    Log1pOp,
339    "math.log1p",
340    "`log(1 + x)`, accurate near zero.",
341    Float
342);
343unary_op!(SqrtOp, "math.sqrt", "Square root.", Float);
344unary_op!(RsqrtOp, "math.rsqrt", "Reciprocal square root.", Float);
345unary_op!(CbrtOp, "math.cbrt", "Cube root.", Float);
346unary_op!(SinOp, "math.sin", "Sine.", Float);
347unary_op!(CosOp, "math.cos", "Cosine.", Float);
348unary_op!(TanOp, "math.tan", "Tangent.", Float);
349unary_op!(AsinOp, "math.asin", "Arcsine.", Float);
350unary_op!(AcosOp, "math.acos", "Arccosine.", Float);
351unary_op!(AtanOp, "math.atan", "Arctangent.", Float);
352unary_op!(SinhOp, "math.sinh", "Hyperbolic sine.", Float);
353unary_op!(CoshOp, "math.cosh", "Hyperbolic cosine.", Float);
354unary_op!(TanhOp, "math.tanh", "Hyperbolic tangent.", Float);
355unary_op!(AsinhOp, "math.asinh", "Inverse hyperbolic sine.", Float);
356unary_op!(AcoshOp, "math.acosh", "Inverse hyperbolic cosine.", Float);
357unary_op!(AtanhOp, "math.atanh", "Inverse hyperbolic tangent.", Float);
358unary_op!(ErfOp, "math.erf", "Error function.", Float);
359unary_op!(ErfcOp, "math.erfc", "Complementary error function.", Float);
360
361binary_op!(
362    Atan2Op,
363    "math.atan2",
364    "Two-argument arctangent of `lhs / rhs`.",
365    Float
366);
367binary_op!(PowFOp, "math.powf", "Floating-point power.", Float);
368binary_op!(
369    CopySignOp,
370    "math.copysign",
371    "Magnitude of `lhs` with the sign of `rhs`.",
372    Float
373);
374
375unary_op!(
376    AbsIOp,
377    "math.absi",
378    "Integer absolute value; the minimum value wraps.",
379    Integer
380);
381unary_op!(
382    CountLeadingZerosOp,
383    "math.ctlz",
384    "Counts leading zero bits.",
385    Integer
386);
387unary_op!(
388    CountTrailingZerosOp,
389    "math.cttz",
390    "Counts trailing zero bits.",
391    Integer
392);
393unary_op!(CtPopOp, "math.ctpop", "Counts set bits.", Integer);
394binary_op!(
395    IPowIOp,
396    "math.ipowi",
397    "Signed integer power; negative exponents follow MLIR semantics.",
398    Integer
399);
400
401classification_op!(IsNaNOp, "math.isnan", "Tests for NaN.");
402classification_op!(
403    IsInfOp,
404    "math.isinf",
405    "Tests for positive or negative infinity."
406);
407classification_op!(IsFiniteOp, "math.isfinite", "Tests for a finite value.");
408classification_op!(
409    IsNormalOp,
410    "math.isnormal",
411    "Tests for a normal (not subnormal) value."
412);
413
414/// Fused multiply-add `a * b + c` with a single rounding.
415///
416/// All operands and the result share one floating-point scalar, tensor, or vector type.
417#[pliron_op(
418    name = "math.fma",
419    format = "$0 `,` $1 `,` $2 attr_dict `:` typesig",
420    interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
421)]
422pub struct FmaOp;
423
424impl FmaOp {
425    /// Allocates the operation; the result type is `a`'s type.
426    pub fn new(ctx: &mut Context, a: Value, b: Value, c: Value) -> Self {
427        let ty = a.get_type(ctx);
428        new_op(ctx, |op| Self { op }, ty, vec![a, b, c])
429    }
430    /// Returns the operand at `index` (`0` = `a`, `1` = `b`, `2` = `c`).
431    pub fn operand(&self, ctx: &Context, index: usize) -> Value {
432        self.get_operation().deref(ctx).get_operand(index)
433    }
434}
435result_accessor!(FmaOp);
436fastmath_accessors!(FmaOp);
437
438impl Verify for FmaOp {
439    fn verify(&self, ctx: &Context) -> Result<()> {
440        verify_same_type(ctx, self.get_operation(), 3, Domain::Float)
441    }
442}
443
444/// Clamps `value` to `[min, max]`: `maxf(minf(value, max), min)`; poison when `min > max`.
445///
446/// All operands and the result share one floating-point scalar, tensor, or vector type.
447#[pliron_op(
448    name = "math.clampf",
449    format = "$0 `,` $1 `,` $2 attr_dict `:` typesig",
450    interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
451)]
452pub struct ClampFOp;
453
454impl ClampFOp {
455    /// Allocates the operation; the result type is `value`'s type.
456    pub fn new(ctx: &mut Context, value: Value, min: Value, max: Value) -> Self {
457        let ty = value.get_type(ctx);
458        new_op(ctx, |op| Self { op }, ty, vec![value, min, max])
459    }
460    /// Returns the clamped operand.
461    pub fn value(&self, ctx: &Context) -> Value {
462        self.get_operation().deref(ctx).get_operand(0)
463    }
464    /// Returns the lower bound.
465    pub fn min(&self, ctx: &Context) -> Value {
466        self.get_operation().deref(ctx).get_operand(1)
467    }
468    /// Returns the upper bound.
469    pub fn max(&self, ctx: &Context) -> Value {
470        self.get_operation().deref(ctx).get_operand(2)
471    }
472}
473result_accessor!(ClampFOp);
474fastmath_accessors!(ClampFOp);
475
476impl Verify for ClampFOp {
477    fn verify(&self, ctx: &Context) -> Result<()> {
478        verify_same_type(ctx, self.get_operation(), 3, Domain::Float)
479    }
480}
481
482/// Computes the sine and cosine of one operand together.
483///
484/// The operand and both results share one floating-point scalar, tensor, or vector type.
485#[pliron_op(
486    name = "math.sincos",
487    format = "$0 attr_dict `:` typesig",
488    interfaces = [pliron::builtin::op_interfaces::NResultsInterface<2>]
489)]
490pub struct SincosOp;
491
492impl SincosOp {
493    /// Allocates the operation; both results take `operand`'s type.
494    pub fn new(ctx: &mut Context, operand: Value) -> Self {
495        let ty = operand.get_type(ctx);
496        Self {
497            op: Operation::new(
498                ctx,
499                Self::get_concrete_op_info(),
500                vec![ty, ty],
501                vec![operand],
502                vec![],
503                0,
504            ),
505        }
506    }
507    /// Returns the operand.
508    pub fn operand(&self, ctx: &Context) -> Value {
509        self.get_operation().deref(ctx).get_operand(0)
510    }
511    /// Returns the sine result.
512    pub fn sin(&self, ctx: &Context) -> Value {
513        self.get_operation().deref(ctx).get_result(0)
514    }
515    /// Returns the cosine result.
516    pub fn cos(&self, ctx: &Context) -> Value {
517        self.get_operation().deref(ctx).get_result(1)
518    }
519}
520fastmath_accessors!(SincosOp);
521
522impl Verify for SincosOp {
523    fn verify(&self, ctx: &Context) -> Result<()> {
524        let op = self.get_operation();
525        let (operands, results, regions, successors) = {
526            let ir = op.deref(ctx);
527            (
528                ir.get_num_operands(),
529                ir.get_num_results(),
530                ir.num_regions(),
531                ir.get_num_successors(),
532            )
533        };
534        if operands != 1 || results != 2 || regions != 0 || successors != 0 {
535            return Err(fail(
536                ctx,
537                op,
538                "expected one operand, two results, and no regions",
539            ));
540        }
541        let ty = operand_type(ctx, op, 0);
542        let ir = op.deref(ctx);
543        if ir.get_result(0).get_type(ctx) != ty || ir.get_result(1).get_type(ctx) != ty {
544            return Err(fail(ctx, op, "operand and results must have the same type"));
545        }
546        drop(ir);
547        let shaped = classify(ctx, op, ty)?;
548        if !in_domain(ctx, shaped.element, Domain::Float) {
549            return Err(fail(ctx, op, "element type must be floating point"));
550        }
551        Ok(())
552    }
553}
554
555/// Floating-point base raised to a signless-integer exponent.
556///
557/// The base and result share one floating-point type; the exponent has the
558/// base's shape with a signless-integer element.
559#[pliron_op(
560    name = "math.fpowi",
561    format = "$0 `,` $1 attr_dict `:` typesig",
562    interfaces = [pliron::builtin::op_interfaces::OneResultInterface]
563)]
564pub struct FPowIOp;
565
566impl FPowIOp {
567    /// Allocates the operation; the result type is the base's type.
568    pub fn new(ctx: &mut Context, base: Value, exponent: Value) -> Self {
569        let ty = base.get_type(ctx);
570        new_op(ctx, |op| Self { op }, ty, vec![base, exponent])
571    }
572    /// Returns the floating-point base.
573    pub fn base(&self, ctx: &Context) -> Value {
574        self.get_operation().deref(ctx).get_operand(0)
575    }
576    /// Returns the integer exponent.
577    pub fn exponent(&self, ctx: &Context) -> Value {
578        self.get_operation().deref(ctx).get_operand(1)
579    }
580}
581result_accessor!(FPowIOp);
582fastmath_accessors!(FPowIOp);
583
584impl Verify for FPowIOp {
585    fn verify(&self, ctx: &Context) -> Result<()> {
586        let op = self.get_operation();
587        verify_structure(ctx, op, 2)?;
588        let result = result_type(ctx, op);
589        if operand_type(ctx, op, 0) != result {
590            return Err(fail(ctx, op, "base and result must have the same type"));
591        }
592        let base = classify(ctx, op, result)?;
593        if !shaped::is_float(ctx, base.element) {
594            return Err(fail(ctx, op, "base must be floating-point"));
595        }
596        let exponent = classify(ctx, op, operand_type(ctx, op, 1))?;
597        if !exponent.same_shape(&base) || shaped::signless_width(ctx, exponent.element).is_none() {
598            return Err(fail(
599                ctx,
600                op,
601                "exponent must be a signless integer with the base's shape",
602            ));
603        }
604        Ok(())
605    }
606}
pliron_math/ops/definitions.rs

pliron_math/ops/
definitions.rs
