ruda_kernel/dsl/frontend/
integer_power.rs

1use ruda_core::ir::{
2    ElemType, FloatKind, Instruction, IntKind, ManagedVariable, Operator, UIntKind, UnaryOperator,
3    Variable,
4};
5
6use crate::dsl::prelude::*;
7
8define_scalar!(Compute);
9define_scalar!(Exponent);
10define_scalar!(Magnitude);
11define_size!(Width);
12
13#[ruda]
14fn integer_power<F: Float, I: Int, U: Int, N: Size>(
15    lhs: Vector<F, N>,
16    rhs: Vector<I, N>,
17    #[comptime] signed: bool,
18) -> Vector<F, N> {
19    let mut output = Vector::<F, N>::from_int(1);
20    #[unroll]
21    for lane in 0..N::value() {
22        let mut base = lhs[lane];
23        let exponent = rhs[lane];
24        let mut magnitude = U::cast_from(exponent);
25        if comptime![signed] {
26            if exponent < I::from_int(0) {
27                magnitude = U::from_int(0) - magnitude;
28                base = F::from_int(1) / base;
29            }
30        }
31        let mut result = F::from_int(1);
32        while magnitude != U::from_int(0) {
33            if (magnitude & U::from_int(1)) != U::from_int(0) {
34                result = result * base;
35            }
36            magnitude = magnitude >> U::from_int(1);
37            if magnitude != U::from_int(0) {
38                base = base * base;
39            }
40        }
41        output[lane] = result;
42    }
43    output
44}
45
46#[allow(missing_docs)]
47pub fn expand_integer_power(scope: &mut Scope, lhs: Variable, rhs: Variable, out: Variable) {
48    let compute = match lhs.elem_type() {
49        ElemType::Float(FloatKind::F16 | FloatKind::BF16) => FloatKind::F32,
50        ElemType::Float(kind @ (FloatKind::F32 | FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F64)) => kind,
51        _ => panic!("integer power lowering requires a supported floating base"),
52    };
53    let (signed, magnitude) = match rhs.elem_type() {
54        ElemType::Int(IntKind::I64) => (true, UIntKind::U64),
55        ElemType::UInt(UIntKind::U64) => (false, UIntKind::U64),
56        ElemType::Int(_) => (true, UIntKind::U32),
57        ElemType::UInt(_) => (false, UIntKind::U32),
58        _ => panic!("integer power lowering requires an integer exponent"),
59    };
60    let width = out.vector_size();
61    let exponent_elem = match rhs.elem_type() {
62        ElemType::Int(IntKind::I8 | IntKind::I16) => ElemType::Int(IntKind::I32),
63        ElemType::UInt(UIntKind::U8 | UIntKind::U16) => ElemType::UInt(UIntKind::U32),
64        elem => elem,
65    };
66    scope.register_type::<Compute>(compute.into());
67    scope.register_type::<Exponent>(exponent_elem.into());
68    scope.register_type::<Magnitude>(magnitude.into());
69    scope.register_size::<Width>(width);
70    let base = scope.create_local(Type::new(compute.into()).with_vector_size(width));
71    scope.register(Instruction::new(Operator::Cast(UnaryOperator { input: lhs }), *base));
72    let exponent = scope.create_local(Type::scalar(exponent_elem).with_vector_size(width));
73    scope.register(Instruction::new(Operator::Cast(UnaryOperator { input: rhs }), *exponent));
74    let result = integer_power::expand::<Compute, Exponent, Magnitude, Width>(
75        scope, base.into(), exponent.into(), signed,
76    );
77    let result: ManagedVariable = result.into();
78    scope.register(Instruction::new(Operator::Cast(UnaryOperator { input: *result }), out));
79}
ruda_kernel/dsl/frontend/integer_power.rs

ruda_kernel/dsl/frontend/
integer_power.rs
