cubecl_cpp/
target.rs

1use core::fmt::Debug;
2
3use cubecl_core::ir::ContextExt;
4use pliron::{context::Context, r#type::Typed};
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub enum Target {
8    Cuda,
9    Hip,
10    Metal,
11}
12
13#[derive(Debug, Clone, Copy, Default)]
14pub struct Shared;
15#[derive(Debug, Clone, Copy, Default)]
16pub struct Cuda;
17#[derive(Debug, Clone, Copy, Default)]
18pub struct Hip;
19#[derive(Debug, Clone, Copy, Default)]
20pub struct Metal;
21
22impl Target {
23    pub fn ty_prefix(&self, ctx: &Context, ty: impl Typed) -> &'static str {
24        if ty.is_half(ctx) {
25            self.half_prefix()
26        } else if ty.is_half2(ctx) {
27            self.half2_prefix()
28        } else {
29            ""
30        }
31    }
32
33    pub fn half_prefix(&self) -> &'static str {
34        match self {
35            Target::Cuda | Target::Hip => "h",
36            Target::Metal => "",
37        }
38    }
39
40    pub fn half2_prefix(&self) -> &'static str {
41        match self {
42            Target::Cuda | Target::Hip => "h2",
43            Target::Metal => "",
44        }
45    }
46}
47
48pub trait CppTarget: Default + Clone + Copy + Debug + Send + Sync + 'static {
49    fn target() -> Target;
50}
51
52impl CppTarget for Cuda {
53    fn target() -> Target {
54        Target::Cuda
55    }
56}
57impl CppTarget for Hip {
58    fn target() -> Target {
59        Target::Hip
60    }
61}
62impl CppTarget for Metal {
63    fn target() -> Target {
64        Target::Metal
65    }
66}
67
68impl CtxTarget for Context {}
69pub trait CtxTarget: ContextExt {
70    fn target(&self) -> Target {
71        *self.aux_ty::<Target>()
72    }
73    fn set_target(&mut self, value: Target) {
74        self.set_aux_ty(value);
75    }
76}
77
78macro_rules! dispatch_target {
79    ($ctx: expr, $expr: expr) => {{
80        use $crate::target::CtxTarget;
81        match $ctx.target() {
82            $crate::target::Target::Cuda => {
83                type Target = $crate::target::Cuda;
84                $expr
85            }
86            $crate::target::Target::Hip => {
87                type Target = $crate::target::Hip;
88                $expr
89            }
90            $crate::target::Target::Metal => {
91                type Target = $crate::target::Metal;
92                $expr
93            }
94        }
95    }};
96}
97pub(crate) use dispatch_target;
98
99use crate::shared::ty::TypedExtCPP;
cubecl_cpp/target.rs

cubecl_cpp/
target.rs
