#[cutile::module]
mod kernels {
use cutile::core::*;
pub fn activation_unary_store<const OP: i32, const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let one = constant(f16::from_f32(1.0), out.shape());
let half = constant(f16::from_f32(0.5), out.shape());
let third = constant(f16::from_f32(0.33333334), out.shape());
let three = constant(f16::from_f32(3.0), out.shape());
let six = constant(f16::from_f32(6.0), out.shape());
let pi = constant(f16::from_f32(std::f32::consts::PI), out.shape());
let one_eighty = constant(f16::from_f32(180.0), out.shape());
let selu_alpha = constant(f16::from_f32(1.6732632), out.shape());
let selu_scale = constant(f16::from_f32(1.050701), out.shape());
let value = if OP == 0 {
zero - x
} else if OP == 1 {
absf(x)
} else if OP == 2 {
sqrt(x, rounding::NearestEven, ftz::Disabled)
} else if OP == 3 {
exp(x)
} else if OP == 4 {
log(x)
} else if OP == 5 {
sin(x)
} else if OP == 6 {
cos(x)
} else if OP == 7 {
tan(x)
} else if OP == 8 {
sinh(x)
} else if OP == 9 {
cosh(x)
} else if OP == 10 {
tanh(x)
} else if OP == 11 {
atan2(x, sqrt(one - x * x, rounding::NearestEven, ftz::Disabled))
} else if OP == 12 {
atan2(sqrt(one - x * x, rounding::NearestEven, ftz::Disabled), x)
} else if OP == 13 {
atan2(x, one)
} else if OP == 14 {
ceil(x)
} else if OP == 15 {
floor(x)
} else if OP == 16 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
select(positive, floor(x), ceil(x))
} else if OP == 17 {
one / x
} else if OP == 18 {
x * x
} else if OP == 19 {
rsqrt(x, ftz::Disabled)
} else if OP == 20 {
log(x + sqrt(x * x + one, rounding::NearestEven, ftz::Disabled))
} else if OP == 21 {
log(x + sqrt(x * x - one, rounding::NearestEven, ftz::Disabled))
} else if OP == 22 {
half * log((one + x) / (one - x))
} else if OP == 23 {
max_tile(x, zero)
} else if OP == 24 {
one / (one + exp(zero - x))
} else if OP == 25 {
x / (one + exp(zero - x))
} else if OP == 26 {
let sqrt_2_over_pi = constant(f16::from_f32(0.7978846), out.shape());
let cubic_coeff = constant(f16::from_f32(0.044715), out.shape());
half * x * (one + tanh(sqrt_2_over_pi * (x + cubic_coeff * x * x * x)))
} else if OP == 27 {
let negative = cmpf(x, zero, predicate::LessThan, cmp_ordering::Ordered);
let positive = cmpf(x, zero, predicate::GreaterThan, cmp_ordering::Ordered);
select(negative, zero - one, select(positive, one, zero))
} else if OP == 28 {
log(one + x)
} else if OP == 29 {
exp(x) - one
} else if OP == 30 {
exp2(x, ftz::Disabled)
} else if OP == 31 {
log2(x)
} else if OP == 32 {
let ln_10 = constant(f16::from_f32(std::f32::consts::LN_10), out.shape());
log(x) / ln_10
} else if OP == 33 {
let root = pow(absf(x), third);
let negative = cmpf(x, zero, predicate::LessThan, cmp_ordering::Ordered);
select(negative, zero - root, root)
} else if OP == 34 {
max_tile(x, zero) + log(one + exp(zero - absf(x)))
} else if OP == 35 {
min_tile(max_tile((x + three) / six, zero), one)
} else if OP == 36 {
x * min_tile(max_tile((x + three) / six, zero), one)
} else if OP == 37 {
x * tanh(max_tile(x, zero) + log(one + exp(zero - absf(x))))
} else if OP == 38 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
selu_scale * select(positive, x, selu_alpha * (exp(x) - one))
} else if OP == 39 {
x / (one + absf(x))
} else if OP == 40 {
x * pi / one_eighty
} else if OP == 41 {
x * one_eighty / pi
} else if OP == 42 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
let truncated = select(positive, floor(x), ceil(x));
x - truncated
} else if OP == 43 {
min_tile(max_tile(x, zero), six)
} else if OP == 44 {
x - tanh(x)
} else if OP == 45 {
zero - (max_tile(zero - x, zero) + log(one + exp(zero - absf(zero - x))))
} else {
x
};
out.store(value);
}
#[cutile::entry()]
pub fn swiglu_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
gate: &Tensor<f16, { [-1] }>,
) {
let input_tile: Tile<f16, { [B] }> = load_tile_like(input, out);
let gate_tile: Tile<f16, { [B] }> = load_tile_like(gate, out);
let input_tile: Tile<f32, { [B] }> = convert_tile(input_tile);
let gate_tile: Tile<f32, { [B] }> = convert_tile(gate_tile);
let one = constant(1.0f32, out.shape());
let zero = constant(0.0f32, out.shape());
let value = input_tile * gate_tile / (one + exp(zero - gate_tile));
out.store(convert_tile(value));
}
#[cutile::entry()]
pub fn relu_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<23i32, B>(out, input);
}
#[cutile::entry()]
pub fn relu6_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<43i32, B>(out, input);
}
#[cutile::entry()]
pub fn sigmoid_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<24i32, B>(out, input);
}
#[cutile::entry()]
pub fn silu_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<25i32, B>(out, input);
}
#[cutile::entry()]
pub fn gelu_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<26i32, B>(out, input);
}
#[cutile::entry()]
pub fn softplus_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<34i32, B>(out, input);
}
#[cutile::entry()]
pub fn hard_sigmoid_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<35i32, B>(out, input);
}
#[cutile::entry()]
pub fn hard_swish_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<36i32, B>(out, input);
}
#[cutile::entry()]
pub fn mish_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<37i32, B>(out, input);
}
#[cutile::entry()]
pub fn selu_f16<const B: i32>(out: &mut Tensor<f16, { [B] }>, input: &Tensor<f16, { [-1] }>) {
activation_unary_store::<38i32, B>(out, input);
}
#[cutile::entry()]
pub fn softsign_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<39i32, B>(out, input);
}
#[cutile::entry()]
pub fn tanhshrink_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<44i32, B>(out, input);
}
#[cutile::entry()]
pub fn log_sigmoid_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
) {
activation_unary_store::<45i32, B>(out, input);
}
#[cutile::entry()]
pub fn leaky_relu_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
negative_slope: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
out.store(select(
positive,
x,
x * broadcast_scalar(negative_slope, out.shape()),
));
}
#[cutile::entry()]
pub fn elu_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
alpha: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let one = constant(f16::from_f32(1.0), out.shape());
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
out.store(select(
positive,
x,
broadcast_scalar(alpha, out.shape()) * (exp(x) - one),
));
}
#[cutile::entry()]
pub fn celu_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
alpha: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let one = constant(f16::from_f32(1.0), out.shape());
let alpha = broadcast_scalar(alpha, out.shape());
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
out.store(select(positive, x, alpha * (exp(x / alpha) - one)));
}
#[cutile::entry()]
pub fn hardshrink_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
lambda: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let lambda = broadcast_scalar(lambda, out.shape());
let keep_positive = cmpf(x, lambda, predicate::GreaterThan, cmp_ordering::Ordered);
let keep_negative = cmpf(x, zero - lambda, predicate::LessThan, cmp_ordering::Ordered);
out.store(select(keep_positive | keep_negative, x, zero));
}
#[cutile::entry()]
pub fn softshrink_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
lambda: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let zero = constant(f16::from_f32(0.0), out.shape());
let lambda = broadcast_scalar(lambda, out.shape());
let positive = cmpf(x, lambda, predicate::GreaterThan, cmp_ordering::Ordered);
let negative = cmpf(x, zero - lambda, predicate::LessThan, cmp_ordering::Ordered);
out.store(select(
positive,
x - lambda,
select(negative, x + lambda, zero),
));
}
#[cutile::entry()]
pub fn threshold_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
threshold: f16,
value: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let keep = cmpf(
x,
broadcast_scalar(threshold, out.shape()),
predicate::GreaterThan,
cmp_ordering::Ordered,
);
out.store(select(keep, x, broadcast_scalar(value, out.shape())));
}
#[cutile::entry()]
pub fn hardtanh_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
min_value: f16,
max_value: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let min_value = broadcast_scalar(min_value, out.shape());
let max_value = broadcast_scalar(max_value, out.shape());
out.store(min_tile(max_tile(x, min_value), max_value));
}
}
pub use kernels::*;