cubecl-llvm 0.11.0-pre.4

LLVM compiler for CubeCL
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

pub mod atomic;
pub mod cmp;
pub mod constant;
pub mod general;
pub mod insert;
pub mod math;
pub mod memory;
pub mod ty;
pub mod vector;

use crate::prelude::*;

pub mod prelude {
    pub use crate::prelude::*;
}

#[type_interface]
pub trait CubeToLLVMType {
    verify_ty_succ!();
    fn convert(&self, ctx: &Context) -> TypeHandle;
}

#[op_interface]
pub trait ToLLVMDialect {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        operands_info: &OperandsInfo,
    ) -> Result<()>;

    fn verify(_op: &dyn Op, _ctx: &Context) -> Result<()>
    where
        Self: Sized,
    {
        Ok(())
    }
}

pub type CubeToLLVMPass = DialectConversionPass<CubeToLLVM>;

#[derive(Default, NamedRewrite)]
pub struct CubeToLLVM;

impl DialectConversion for CubeToLLVM {
    fn can_convert_op(&self, ctx: &Context, op: Ptr<Operation>) -> bool {
        op_impls::<dyn ToLLVMDialect>(&*Operation::get_op_dyn(op, ctx))
    }

    fn can_convert_type(&self, ctx: &Context, ty: TypeHandle) -> bool {
        type_cast::<dyn CubeToLLVMType>(&*ty.deref(ctx)).is_some()
    }

    fn convert_type(&mut self, ctx: &mut Context, ty: TypeHandle) -> Result<TypeHandle> {
        let ty = ty.deref(ctx);
        let ty = type_cast::<dyn CubeToLLVMType>(&*ty).unwrap();
        let ty = ty.convert(ctx);
        Ok(ty)
    }

    fn rewrite(
        &mut self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        op: Ptr<Operation>,
        operands_info: &OperandsInfo,
    ) -> Result<()> {
        let op_dyn = Operation::get_op_dyn(op, ctx);
        let to_llvm_op = op_cast::<dyn ToLLVMDialect>(&*op_dyn)
            .expect("Matched Op must implement ToLLVMDialect");
        to_llvm_op.rewrite(ctx, rewriter, operands_info)
    }
}