rtx-backend-metal: align to current rtx-metal API
rtx-metal dropped tensor_ops::{sin,cos,pow,clamp,gt_scalar,var} and never
had nn::{max_pool2d,avg_pool2d}; nn::conv2d grew a scalar 14-arg signature.
Give the missing ops correct host fallbacks (the sum_dim pattern), and
dispatch conv2d to the Metal kernel when its restricted signature applies
(symmetric stride/padding, dilation 1, groups 1), host fallback otherwise.
cargo test -p rtx-backend-metal: 22/22 parity tests pass on-device (M5).
Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
quantum
co-authored by
Claude Fable 5
Claude Fable 5
parent
4e90177aa9
commit
72b41e3167
@@ -1,6 +1,10 @@
|
||||
//! Convolution and pooling operations using Metal GPU shaders.
|
||||
//! Convolution and pooling operations.
|
||||
//!
|
||||
//! These operations dispatch to Metal compute kernels for GPU-accelerated execution.
|
||||
//! `conv2d` dispatches to the Metal `conv2d_f32` kernel when the parameters
|
||||
//! fit its signature (symmetric stride/padding, dilation 1, groups 1); all
|
||||
//! other cases — and both pooling ops, which have no Metal kernel in
|
||||
//! rtx-metal — are computed on host (same fallback pattern as
|
||||
//! `reduction::sum_dim`).
|
||||
|
||||
use crate::MetalTensorPrimitive;
|
||||
use rtx_metal::{MetalBuffer, MetalBufferUsage};
|
||||
@@ -8,7 +12,7 @@ use rtx_metal::{MetalBuffer, MetalBufferUsage};
|
||||
#[cfg(target_os = "macos")]
|
||||
use rtx_metal::ops::nn;
|
||||
|
||||
/// 2D Convolution using Metal GPU kernel.
|
||||
/// 2D Convolution.
|
||||
pub fn conv2d(
|
||||
input: &MetalTensorPrimitive<4>,
|
||||
weight: &MetalTensorPrimitive<4>,
|
||||
@@ -18,8 +22,8 @@ pub fn conv2d(
|
||||
dilation: [usize; 2],
|
||||
groups: usize,
|
||||
) -> MetalTensorPrimitive<4> {
|
||||
let [batch, _in_channels, in_h, in_w] = input.shape;
|
||||
let [out_channels, _in_channels_per_group, kernel_h, kernel_w] = weight.shape;
|
||||
let [batch, in_channels, in_h, in_w] = input.shape;
|
||||
let [out_channels, in_channels_per_group, kernel_h, kernel_w] = weight.shape;
|
||||
|
||||
let out_h = (in_h + 2 * padding[0] - dilation[0] * (kernel_h - 1) - 1) / stride[0] + 1;
|
||||
let out_w = (in_w + 2 * padding[1] - dilation[1] * (kernel_w - 1) - 1) / stride[1] + 1;
|
||||
@@ -28,40 +32,93 @@ pub fn conv2d(
|
||||
let out_numel: usize = out_shape.iter().product();
|
||||
|
||||
let device = input.device.metal_device();
|
||||
let mut output = MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for conv2d");
|
||||
|
||||
// The Metal kernel only supports symmetric stride/padding, dilation 1,
|
||||
// groups 1. Dispatch to it when possible.
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
nn::conv2d(
|
||||
device,
|
||||
input.data(),
|
||||
weight.data(),
|
||||
bias.map(|b| b.data()),
|
||||
&mut output,
|
||||
input.shape,
|
||||
weight.shape,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
)
|
||||
.expect("Failed to execute Metal conv2d kernel");
|
||||
if stride[0] == stride[1]
|
||||
&& padding[0] == padding[1]
|
||||
&& dilation == [1, 1]
|
||||
&& groups == 1
|
||||
{
|
||||
let mut output = MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for conv2d");
|
||||
nn::conv2d(
|
||||
device,
|
||||
input.data(),
|
||||
weight.data(),
|
||||
bias.map(|b| b.data()),
|
||||
&mut output,
|
||||
batch,
|
||||
in_channels,
|
||||
out_channels,
|
||||
in_h,
|
||||
in_w,
|
||||
kernel_h,
|
||||
kernel_w,
|
||||
stride[0],
|
||||
padding[0],
|
||||
)
|
||||
.expect("Failed to execute Metal conv2d kernel");
|
||||
return MetalTensorPrimitive::new(output, out_shape, input.device.clone());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let _input_data = input.to_vec();
|
||||
let _weight_data = weight.to_vec();
|
||||
let _bias_data = bias.map(|b| b.to_vec());
|
||||
let result = vec![0.0f32; out_numel];
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
// Host fallback: naive direct convolution with full stride / padding /
|
||||
// dilation / groups support.
|
||||
let input_data = input.to_vec();
|
||||
let weight_data = weight.to_vec();
|
||||
let bias_data = bias.map(|b| b.to_vec());
|
||||
let mut result = vec![0.0f32; out_numel];
|
||||
|
||||
let out_channels_per_group = out_channels / groups;
|
||||
|
||||
for b in 0..batch {
|
||||
for oc in 0..out_channels {
|
||||
let group = oc / out_channels_per_group;
|
||||
for oh in 0..out_h {
|
||||
for ow in 0..out_w {
|
||||
let mut acc = bias_data.as_ref().map_or(0.0, |bd| bd[oc]);
|
||||
for ic in 0..in_channels_per_group {
|
||||
let ic_global = group * in_channels_per_group + ic;
|
||||
for kh in 0..kernel_h {
|
||||
for kw in 0..kernel_w {
|
||||
let ih = (oh * stride[0] + kh * dilation[0]) as isize
|
||||
- padding[0] as isize;
|
||||
let iw = (ow * stride[1] + kw * dilation[1]) as isize
|
||||
- padding[1] as isize;
|
||||
if ih >= 0
|
||||
&& iw >= 0
|
||||
&& (ih as usize) < in_h
|
||||
&& (iw as usize) < in_w
|
||||
{
|
||||
let in_idx = b * in_channels * in_h * in_w
|
||||
+ ic_global * in_h * in_w
|
||||
+ ih as usize * in_w
|
||||
+ iw as usize;
|
||||
let w_idx = oc * in_channels_per_group * kernel_h * kernel_w
|
||||
+ ic * kernel_h * kernel_w
|
||||
+ kh * kernel_w
|
||||
+ kw;
|
||||
acc += input_data[in_idx] * weight_data[w_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let out_idx =
|
||||
b * out_channels * out_h * out_w + oc * out_h * out_w + oh * out_w + ow;
|
||||
result[out_idx] = acc;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, out_shape, input.device.clone())
|
||||
}
|
||||
|
||||
/// 2D max pooling using Metal GPU kernel.
|
||||
/// 2D max pooling (host fallback — no Metal kernel in rtx-metal).
|
||||
pub fn max_pool2d(
|
||||
input: &MetalTensorPrimitive<4>,
|
||||
kernel_size: [usize; 2],
|
||||
@@ -77,67 +134,40 @@ pub fn max_pool2d(
|
||||
let out_numel: usize = out_shape.iter().product();
|
||||
|
||||
let device = input.device.metal_device();
|
||||
let mut output = MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for max_pool2d");
|
||||
let input_data = input.to_vec();
|
||||
let mut result = vec![f32::NEG_INFINITY; out_numel];
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
nn::max_pool2d(
|
||||
device,
|
||||
input.data(),
|
||||
&mut output,
|
||||
input.shape,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
)
|
||||
.expect("Failed to execute Metal max_pool2d kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let input_data = input.to_vec();
|
||||
let mut result = vec![f32::NEG_INFINITY; out_numel];
|
||||
|
||||
for b in 0..batch {
|
||||
for c in 0..channels {
|
||||
for oh in 0..out_h {
|
||||
for ow in 0..out_w {
|
||||
let mut max_val = f32::NEG_INFINITY;
|
||||
for kh in 0..kernel_size[0] {
|
||||
for kw in 0..kernel_size[1] {
|
||||
let ih = oh * stride[0] + kh;
|
||||
let iw = ow * stride[1] + kw;
|
||||
if ih >= padding[0]
|
||||
&& iw >= padding[1]
|
||||
&& ih < in_h + padding[0]
|
||||
&& iw < in_w + padding[1]
|
||||
{
|
||||
let ih_actual = ih - padding[0];
|
||||
let iw_actual = iw - padding[1];
|
||||
let in_idx = b * channels * in_h * in_w
|
||||
+ c * in_h * in_w
|
||||
+ ih_actual * in_w
|
||||
+ iw_actual;
|
||||
max_val = max_val.max(input_data[in_idx]);
|
||||
}
|
||||
for b in 0..batch {
|
||||
for c in 0..channels {
|
||||
for oh in 0..out_h {
|
||||
for ow in 0..out_w {
|
||||
let mut max_val = f32::NEG_INFINITY;
|
||||
for kh in 0..kernel_size[0] {
|
||||
for kw in 0..kernel_size[1] {
|
||||
let ih = (oh * stride[0] + kh) as isize - padding[0] as isize;
|
||||
let iw = (ow * stride[1] + kw) as isize - padding[1] as isize;
|
||||
if ih >= 0 && iw >= 0 && (ih as usize) < in_h && (iw as usize) < in_w {
|
||||
let in_idx = b * channels * in_h * in_w
|
||||
+ c * in_h * in_w
|
||||
+ ih as usize * in_w
|
||||
+ iw as usize;
|
||||
max_val = max_val.max(input_data[in_idx]);
|
||||
}
|
||||
}
|
||||
let out_idx =
|
||||
b * channels * out_h * out_w + c * out_h * out_w + oh * out_w + ow;
|
||||
result[out_idx] = max_val;
|
||||
}
|
||||
let out_idx =
|
||||
b * channels * out_h * out_w + c * out_h * out_w + oh * out_w + ow;
|
||||
result[out_idx] = max_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, out_shape, input.device.clone())
|
||||
}
|
||||
|
||||
/// 2D average pooling using Metal GPU kernel.
|
||||
/// 2D average pooling (host fallback — no Metal kernel in rtx-metal).
|
||||
pub fn avg_pool2d(
|
||||
input: &MetalTensorPrimitive<4>,
|
||||
kernel_size: [usize; 2],
|
||||
@@ -154,67 +184,39 @@ pub fn avg_pool2d(
|
||||
let out_numel: usize = out_shape.iter().product();
|
||||
|
||||
let device = input.device.metal_device();
|
||||
let mut output = MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for avg_pool2d");
|
||||
let input_data = input.to_vec();
|
||||
let mut result = vec![0.0f32; out_numel];
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
nn::avg_pool2d(
|
||||
device,
|
||||
input.data(),
|
||||
&mut output,
|
||||
input.shape,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
count_include_pad,
|
||||
)
|
||||
.expect("Failed to execute Metal avg_pool2d kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let input_data = input.to_vec();
|
||||
let mut result = vec![0.0f32; out_numel];
|
||||
|
||||
for b in 0..batch {
|
||||
for c in 0..channels {
|
||||
for oh in 0..out_h {
|
||||
for ow in 0..out_w {
|
||||
let mut sum = 0.0f32;
|
||||
let mut count = 0;
|
||||
for kh in 0..kernel_size[0] {
|
||||
for kw in 0..kernel_size[1] {
|
||||
let ih = oh * stride[0] + kh;
|
||||
let iw = ow * stride[1] + kw;
|
||||
if ih >= padding[0]
|
||||
&& iw >= padding[1]
|
||||
&& ih < in_h + padding[0]
|
||||
&& iw < in_w + padding[1]
|
||||
{
|
||||
let ih_actual = ih - padding[0];
|
||||
let iw_actual = iw - padding[1];
|
||||
let in_idx = b * channels * in_h * in_w
|
||||
+ c * in_h * in_w
|
||||
+ ih_actual * in_w
|
||||
+ iw_actual;
|
||||
sum += input_data[in_idx];
|
||||
count += 1;
|
||||
} else if count_include_pad {
|
||||
count += 1;
|
||||
}
|
||||
for b in 0..batch {
|
||||
for c in 0..channels {
|
||||
for oh in 0..out_h {
|
||||
for ow in 0..out_w {
|
||||
let mut sum = 0.0f32;
|
||||
let mut count = 0usize;
|
||||
for kh in 0..kernel_size[0] {
|
||||
for kw in 0..kernel_size[1] {
|
||||
let ih = (oh * stride[0] + kh) as isize - padding[0] as isize;
|
||||
let iw = (ow * stride[1] + kw) as isize - padding[1] as isize;
|
||||
if ih >= 0 && iw >= 0 && (ih as usize) < in_h && (iw as usize) < in_w {
|
||||
let in_idx = b * channels * in_h * in_w
|
||||
+ c * in_h * in_w
|
||||
+ ih as usize * in_w
|
||||
+ iw as usize;
|
||||
sum += input_data[in_idx];
|
||||
count += 1;
|
||||
} else if count_include_pad {
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
let out_idx =
|
||||
b * channels * out_h * out_w + c * out_h * out_w + oh * out_w + ow;
|
||||
result[out_idx] = if count > 0 { sum / count as f32 } else { 0.0 };
|
||||
}
|
||||
let out_idx =
|
||||
b * channels * out_h * out_w + c * out_h * out_w + oh * out_w + ow;
|
||||
result[out_idx] = if count > 0 { sum / count as f32 } else { 0.0 };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, out_shape, input.device.clone())
|
||||
}
|
||||
|
||||
@@ -157,35 +157,24 @@ pub fn min<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimi
|
||||
MetalTensorPrimitive::new(output, [1], tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Variance of all elements using Metal GPU kernel.
|
||||
/// Variance of all elements (population variance, divisor `n`).
|
||||
///
|
||||
/// rtx-metal has no dedicated variance kernel; computed on host
|
||||
/// (same fallback pattern as `sum_dim`).
|
||||
pub fn var<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimitive<1> {
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, 1, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for var");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::var(device, tensor.data(), &mut output)
|
||||
.expect("Failed to execute Metal var kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let n = data.len() as f32;
|
||||
let mean: f32 = data.iter().sum::<f32>() / n;
|
||||
let variance: f32 = data
|
||||
.iter()
|
||||
.map(|&x| {
|
||||
let diff = x - mean;
|
||||
diff * diff
|
||||
})
|
||||
.sum::<f32>()
|
||||
/ n;
|
||||
output = MetalBuffer::from_slice(device, &[variance]).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let n = data.len() as f32;
|
||||
let mean: f32 = data.iter().sum::<f32>() / n;
|
||||
let variance: f32 = data
|
||||
.iter()
|
||||
.map(|&x| {
|
||||
let diff = x - mean;
|
||||
diff * diff
|
||||
})
|
||||
.sum::<f32>()
|
||||
/ n;
|
||||
let output = MetalBuffer::from_slice(device, &[variance]).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, [1], tensor.device.clone())
|
||||
}
|
||||
|
||||
|
||||
@@ -131,132 +131,63 @@ pub fn abs<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimi
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Element-wise sine using Metal GPU kernel.
|
||||
/// Element-wise sine.
|
||||
///
|
||||
/// rtx-metal has no dedicated kernel for this op; computed on host
|
||||
/// (same fallback pattern as `reduction::sum_dim`).
|
||||
pub fn sin<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimitive<D> {
|
||||
let numel = tensor.numel();
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for sin");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::sin(device, tensor.data(), &mut output)
|
||||
.expect("Failed to execute Metal sin kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.sin()).collect();
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.sin()).collect();
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Element-wise cosine using Metal GPU kernel.
|
||||
/// Element-wise cosine.
|
||||
///
|
||||
/// rtx-metal has no dedicated kernel for this op; computed on host
|
||||
/// (same fallback pattern as `reduction::sum_dim`).
|
||||
pub fn cos<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimitive<D> {
|
||||
let numel = tensor.numel();
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for cos");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::cos(device, tensor.data(), &mut output)
|
||||
.expect("Failed to execute Metal cos kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.cos()).collect();
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.cos()).collect();
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Element-wise power using Metal GPU kernel.
|
||||
/// Element-wise power.
|
||||
///
|
||||
/// rtx-metal has no dedicated kernel for this op; computed on host
|
||||
/// (same fallback pattern as `reduction::sum_dim`).
|
||||
pub fn pow<const D: usize>(tensor: &MetalTensorPrimitive<D>, exp: f32) -> MetalTensorPrimitive<D> {
|
||||
let numel = tensor.numel();
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for pow");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::pow(device, tensor.data(), &mut output, exp)
|
||||
.expect("Failed to execute Metal pow kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.powf(exp)).collect();
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.powf(exp)).collect();
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Clamp tensor values to a range using Metal GPU kernel.
|
||||
pub fn clamp<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
min: f32,
|
||||
max: f32,
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
let numel = tensor.numel();
|
||||
/// Clamp tensor values to a range.
|
||||
///
|
||||
/// rtx-metal has no dedicated kernel for this op; computed on host
|
||||
/// (same fallback pattern as `reduction::sum_dim`).
|
||||
pub fn clamp<const D: usize>(tensor: &MetalTensorPrimitive<D>, min: f32,
|
||||
max: f32) -> MetalTensorPrimitive<D> {
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for clamp");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::clamp(device, tensor.data(), &mut output, min, max)
|
||||
.expect("Failed to execute Metal clamp kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.clamp(min, max)).collect();
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|x| x.clamp(min, max)).collect();
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Greater than scalar comparison using Metal GPU kernel.
|
||||
pub fn gt_scalar<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
value: f32,
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
let numel = tensor.numel();
|
||||
/// Greater-than-scalar comparison (1.0 / 0.0 mask).
|
||||
///
|
||||
/// rtx-metal has no dedicated kernel for this op; computed on host
|
||||
/// (same fallback pattern as `reduction::sum_dim`).
|
||||
pub fn gt_scalar<const D: usize>(tensor: &MetalTensorPrimitive<D>, value: f32) -> MetalTensorPrimitive<D> {
|
||||
let device = tensor.device.metal_device();
|
||||
|
||||
let mut output = MetalBuffer::new(device, numel, MetalBufferUsage::Shared)
|
||||
.expect("Failed to allocate output buffer for gt_scalar");
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
tensor_ops::gt_scalar(device, tensor.data(), &mut output, value)
|
||||
.expect("Failed to execute Metal gt_scalar kernel");
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "macos"))]
|
||||
{
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data
|
||||
.iter()
|
||||
.map(|&x| if x > value { 1.0 } else { 0.0 })
|
||||
.collect();
|
||||
output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
}
|
||||
|
||||
let data = tensor.to_vec();
|
||||
let result: Vec<f32> = data.iter().map(|&x| if x > value { 1.0 } else { 0.0 }).collect();
|
||||
let output = MetalBuffer::from_slice(device, &result).expect("Failed to create buffer");
|
||||
MetalTensorPrimitive::new(output, tensor.shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user