cubecl_cpp/cuda/
binary.rs

1use cubecl_core::ir::{
2    dialect::math::{Dp4aOp, SaturatingSAddOp, SaturatingSSubOp},
3    interfaces::TypedExt,
4    prelude::*,
5};
6use cubecl_core::{frontend::polyfills::expand_dp4a_polyfill, ir::Scope};
7
8use crate::{
9    cuda::{cuda_op_with_out, ptx_with_out},
10    shared::{CompilationOptions, lowering::LowerOp},
11    target::Cuda,
12};
13
14cuda_op_with_out!(Dp4aOp, |op, ctx| {
15    let a = op.a(ctx).name(ctx);
16    let b = op.b(ctx).name(ctx);
17    let c = op.c(ctx).name(ctx);
18    format!("__dp4a({a}, {b}, {c})")
19});
20
21#[op_interface_impl]
22impl LowerOp<Cuda> for Dp4aOp {
23    fn should_lower(&self, ctx: &Context) -> bool {
24        !ctx.aux_ty::<CompilationOptions>().supports_features.dp4a
25    }
26
27    fn lower(&self, scope: &Scope) -> Vec<Value> {
28        let ctx = scope.ctx();
29        vec![expand_dp4a_polyfill(
30            scope,
31            self.a(ctx),
32            self.b(ctx),
33            self.c(ctx),
34        )]
35    }
36}
37
38ptx_with_out!(
39    SaturatingSAddOp,
40    |_, _| "add.sat.s32 $0, $1, $2;".into(),
41    |op, ctx| op.result_type(ctx).is_int_of_width(ctx, 32)
42        && op.result_type(ctx).is_signed_int(ctx)
43);
44ptx_with_out!(
45    SaturatingSSubOp,
46    |_, _| "sub.sat.s32 $0, $1, $2;".into(),
47    |op, ctx| op.result_type(ctx).is_int_of_width(ctx, 32)
48        && op.result_type(ctx).is_signed_int(ctx)
49);
cubecl_cpp/cuda/binary.rs

cubecl_cpp/cuda/
binary.rs
