cubecl_cpp/metal/
vector.rs

1use cubecl_core::ir::{dialect::vector::*, interfaces::TypedExt, prelude::*};
2
3use crate::{
4    metal::metal_op_with_out,
5    shared::{scoped_block, ty::TypeExtCPP},
6};
7
8metal_op_with_out!(MagnitudeOp, |op, ctx| {
9    let input = op.input(ctx).name(ctx);
10    let scalar_ty = op.result_type(ctx).to_cpp(ctx);
11    let vec = op.input(ctx).vector_size(ctx);
12    let input = format!("reinterpret_cast<const thread {scalar_ty}{vec}&>({input})");
13    format!("length({input})")
14});
15
16metal_op_with_out!(NormalizeOp, |op, ctx| {
17    let input = op.input(ctx).name(ctx);
18    let scalar_ty = op.result_type(ctx).to_cpp(ctx);
19    let vec = op.input(ctx).vector_size(ctx);
20    let input = format!("reinterpret_cast<const thread {scalar_ty}{vec}&>({input})");
21    format!("normalize({input})")
22});
23
24metal_op_with_out!(FDotOp, |op, ctx| {
25    let lhs = op.lhs(ctx).name(ctx);
26    let rhs = op.rhs(ctx).name(ctx);
27    let scalar_ty = op.result_type(ctx).to_cpp(ctx);
28    let vec = op.lhs(ctx).vector_size(ctx);
29    let reinterpret = format!("reinterpret_cast<const thread {scalar_ty}{vec}&>");
30    format!("dot({reinterpret}({lhs}), {reinterpret}({rhs}))")
31});
32
33// Workaround for Metal compiler bug that causes combinatorial explosion with large aggregate
34// literal chains. No idea what they're doing wrong but this works around it. Keep it Metal only
35// because it's slightly less clean and less analyzable than the literal constructor.
36metal_op_with_out!(CompositeInsertOp, |op, ctx| {
37    assert!(op.composite(ctx).is_vector(ctx));
38    let vector = op.composite(ctx).name(ctx);
39    let value = op.value(ctx).name(ctx);
40    let index = op.index(ctx).0;
41    let vector_ty = op.composite(ctx).get_type(ctx).to_cpp(ctx);
42    scoped_block!(
43        format!("{vector_ty} tmp = {vector};")
44        format!("tmp.i_{index} = {value};")
45        "return tmp;"
46    )
47});
48
49metal_op_with_out!(VectorInsertDynamicOp, |op, ctx| {
50    let vector = op.vector(ctx).name(ctx);
51    let value = op.value(ctx).name(ctx);
52    let index = op.index(ctx).name(ctx);
53    let elem_ty = op.value(ctx).get_type(ctx).to_cpp(ctx);
54    let vector_ty = op.vector(ctx).get_type(ctx).to_cpp(ctx);
55    scoped_block!(
56        format!("{vector_ty} tmp = {vector};")
57        format!("*(reinterpret_cast<thread {elem_ty}*>(&tmp) + {index}) = {value};")
58        "return tmp;"
59    )
60});
61
62metal_op_with_out!(VectorExtractDynamicOp, |op, ctx| {
63    let vector = op.vector(ctx).name(ctx);
64    let index = op.index(ctx).name(ctx);
65    let elem_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
66    format!("reinterpret_cast<const thread {elem_ty}*>(&{vector})[{index}]")
67});
cubecl_cpp/metal/vector.rs

cubecl_cpp/metal/
vector.rs
