cubecl_cpp/shared/
convert.rs

1//! Convert unsupported types ahead of time. Also convert auto-promoted types manually so we can
2//! properly preserve the semantics of the actual code in IR.
3
4use cubecl_core::ir::{
5    NamedRewrite,
6    dialect::{OperationPtrExt, general::CastOp},
7    interfaces::ValueExt,
8    prelude::{Context, MatchRewriter, Operation, Ptr, Result},
9    rewrite::MatchRewritePass,
10    types::scalar::Float32Type,
11    verify_op_succ,
12};
13use pliron::{
14    builtin::{
15        op_interfaces::OneResultInterface,
16        types::{IntegerType, Signedness},
17    },
18    derive::op_interface,
19    irbuild::match_rewrite::MatchRewrite,
20    op::Op,
21    r#type::{TypeHandle, Typed},
22    value::Value,
23};
24
25#[op_interface]
26pub trait HalfPromotedOp: OneResultInterface {
27    verify_op_succ!();
28}
29
30#[op_interface]
31pub trait IntPromotedOp: OneResultInterface {
32    verify_op_succ!();
33}
34
35macro_rules! no_half {
36    ($ty: ty) => {
37        #[pliron::derive::op_interface_impl]
38        impl $crate::shared::convert::HalfPromotedOp for $ty {}
39    };
40}
41
42pub(crate) use no_half;
43
44macro_rules! promotes_int {
45    ($ty: ty) => {
46        #[pliron::derive::op_interface_impl]
47        impl $crate::shared::convert::IntPromotedOp for $ty {}
48    };
49}
50pub(crate) use promotes_int;
51
52use crate::shared::ty::TypedExtCPP;
53
54pub type PromoteUnsupportedTypesPass = MatchRewritePass<PromoteUnsupportedTypes>;
55
56#[derive(Default, NamedRewrite)]
57pub struct PromoteUnsupportedTypes;
58
59impl MatchRewrite for PromoteUnsupportedTypes {
60    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
61        let promote_half = op.impls::<dyn HalfPromotedOp>(ctx) && op.result(ctx).is_half(ctx);
62        let promote_int = op.impls::<dyn IntPromotedOp>(ctx) && op.result(ctx).is_small_int(ctx);
63        promote_half || promote_int
64    }
65
66    fn rewrite(
67        &mut self,
68        ctx: &mut Context,
69        _rewriter: &mut MatchRewriter,
70        op: Ptr<Operation>,
71    ) -> Result<()> {
72        if op.impls::<dyn HalfPromotedOp>(ctx) && op.result(ctx).is_half(ctx) {
73            let f32 = Float32Type::get(ctx).to_handle();
74            promote(ctx, op, |ctx, value| value.is_half(ctx), f32);
75        }
76        if op.impls::<dyn IntPromotedOp>(ctx) && op.result(ctx).is_small_signed_int(ctx) {
77            let i32 = IntegerType::get(ctx, 32, Signedness::Signed).to_handle();
78            promote(ctx, op, |ctx, value| value.is_small_int(ctx), i32);
79        }
80        if op.impls::<dyn IntPromotedOp>(ctx) && op.result(ctx).is_small_unsigned_int(ctx) {
81            let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
82            promote(ctx, op, |ctx, value| value.is_small_int(ctx), u32);
83        }
84
85        Ok(())
86    }
87}
88
89fn promote(
90    ctx: &mut Context,
91    op: Ptr<Operation>,
92    pred: impl Fn(&Context, Value) -> bool,
93    promoted_type: TypeHandle,
94) {
95    let res = op.result(ctx);
96    let res_ty = res.get_type(ctx);
97
98    for r#use in op.operands_as_uses(ctx) {
99        let value = r#use.get_def(ctx);
100        if pred(ctx, value) {
101            let cast = CastOp::new(ctx, promoted_type, value);
102            cast.get_operation().insert_before(ctx, op);
103            value.replace_use_with(ctx, r#use, &cast.get_result(ctx));
104        }
105    }
106
107    res.set_type(ctx, promoted_type);
108    let cast_res = CastOp::new(ctx, res_ty, res);
109    let new_res = cast_res.get_result(ctx);
110    cast_res.get_operation().insert_after(ctx, op);
111    res.replace_all_uses_except_with(ctx, cast_res.input_as_use(ctx), &new_res);
112}
cubecl_cpp/shared/convert.rs

cubecl_cpp/shared/
convert.rs
