Skip to content

[MPS] Implement hardshrink metal kernel #155304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
manuelcandales wants to merge 13 commits into from
Closed

[MPS] Implement hardshrink metal kernel #155304

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
f39e761
Update
manuelcandales Jun 6, 2025
f5676f4
Update
manuelcandales Jun 6, 2025
e393dba
Update
manuelcandales Jun 6, 2025
0ac7f8e
Update
manuelcandales Jun 6, 2025
b474056
Update
manuelcandales Jun 6, 2025
f2e018b
Update
manuelcandales Jun 7, 2025
2fdce93
Update
manuelcandales Jun 9, 2025
c635f81
Update
manuelcandales Jun 9, 2025
3a26396
Update
manuelcandales Jun 9, 2025
d48911b
Update
manuelcandales Jun 9, 2025
fcef81e
Update
manuelcandales Jun 9, 2025
5314497
Update
manuelcandales Jun 9, 2025
525ae01
Update
manuelcandales Jun 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions aten/src/ATen/native/mps/MetalShaderLibrary.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,11 +137,13 @@ class MetalShaderLibrary {
void exec_unary_kernel(
TensorIteratorBase& iter,
const std::string& name,
std::optional<int64_t> extra = std::nullopt);
const std::optional<c10::Scalar> alpha = std::nullopt,
const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);
Comment on lines +140 to +141
Copy link
Contributor

@malfet malfet Jun 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This API looks wrong, it should be something like

Suggested change
const std::optional<c10::Scalar> alpha = std::nullopt,
const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);
const std::optional<std::pair<c10::Scalar, c10::ScalarType>> alpha = std::nullopt,

void exec_binary_kernel(
TensorIteratorBase& iter,
const std::string& name,
const std::optional<c10::Scalar> alpha = std::nullopt);
const std::optional<c10::Scalar> alpha = std::nullopt,
const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);

protected:
virtual MTLLibrary_t getLibrary();
Expand Down
31 changes: 19 additions & 12 deletions aten/src/ATen/native/mps/OperationUtils.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -990,11 +990,12 @@ static dispatch_data_t getSectionData(const std::string& name) {

void MetalShaderLibrary::exec_unary_kernel(TensorIteratorBase& iter,
const std::string& name,
std::optional<int64_t> extra) {
std::optional<c10::Scalar> alpha,
std::optional<c10::ScalarType> scalar_arg_type) {
// Decompose 64-bit tensor into 32-bit ones
if (!iter.can_use_32bit_indexing()) {
for (auto&& sub_iter : iter.with_32bit_indexing()) {
exec_unary_kernel(sub_iter, name, extra);
exec_unary_kernel(sub_iter, name, alpha, scalar_arg_type);
}
return;
}
Expand All @@ -1006,11 +1007,13 @@ static dispatch_data_t getSectionData(const std::string& name) {
return;
}
using namespace mps;
auto kernel_name = fmt::format("{}_{}_{}_{}",
const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype();
auto kernel_name = fmt::format("{}_{}_{}_{}{}",
name,
iter.is_contiguous() ? "dense" : "strided",
scalarToMetalTypeString(outputTensor),
scalarToMetalTypeString(inputTensor));
scalarToMetalTypeString(inputTensor),
alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : "");
@autoreleasepool {
auto cplState = getPipelineStateForFunc(kernel_name);

Expand All @@ -1029,8 +1032,8 @@ static dispatch_data_t getSectionData(const std::string& name) {
outputTensor.strides(),
inputTensor.ndimension());
}
if (extra) {
mtl_setBytes(computeEncoder, *extra, iter.is_contiguous() ? 2 : 6);
if (alpha) {
mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), iter.is_contiguous() ? 2 : 6);
}
mtl_dispatch1DJob(computeEncoder, cplState, length);

Expand All @@ -1041,7 +1044,8 @@ static dispatch_data_t getSectionData(const std::string& name) {

void MetalShaderLibrary::exec_binary_kernel(TensorIteratorBase& iter,
const std::string& name,
std::optional<c10::Scalar> alpha) {
std::optional<c10::Scalar> alpha,
std::optional<c10::ScalarType> scalar_arg_type) {
// TODO: Figure a better place to downcast double scalars (probably in tensor iterator itself?)
// Right now running something like 1.0-torch.rand(5, device='mps') will create iterator with
// double as common dtype (because Python floating point are always 64-bit values)
Expand All @@ -1055,7 +1059,7 @@ static dispatch_data_t getSectionData(const std::string& name) {
// Decompose 64-bit tensor into 32-bit ones
if (!iter.can_use_32bit_indexing()) {
for (auto&& sub_iter : iter.with_32bit_indexing()) {
exec_binary_kernel(sub_iter, name, alpha);
exec_binary_kernel(sub_iter, name, alpha, scalar_arg_type);
}
return;
}
Expand All @@ -1081,10 +1085,13 @@ static dispatch_data_t getSectionData(const std::string& name) {
MPSStream* mpsStream = getCurrentMPSStream();
const auto cast_needed = input.scalar_type() != other.scalar_type();
const auto suffix = iter.is_contiguous() ? "dense" : "strided";
const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype();
const auto alpha_suffix = alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : "";
// TODO: Implicitly pass both input and output types to non-cast kernels
const auto kernel_name = cast_needed
? fmt::format("{}_{}_cast_{}", name, suffix, scalarToMetalTypeString(out))
: fmt::format("{}_{}_{}_{}", name, suffix, scalarToMetalTypeString(out), scalarToMetalTypeString(input));
? fmt::format("{}_{}_cast_{}{}", name, suffix, scalarToMetalTypeString(out), alpha_suffix)
: fmt::format(
"{}_{}_{}_{}{}", name, suffix, scalarToMetalTypeString(out), scalarToMetalTypeString(input), alpha_suffix);
dispatch_sync_with_rethrow(mpsStream->queue(), ^() {
@autoreleasepool {
auto computeEncoder = mpsStream->commandEncoder();
Expand All @@ -1098,7 +1105,7 @@ static dispatch_data_t getSectionData(const std::string& name) {
// i.e. it's true for both row-first and column-first tensors
if (iter.is_contiguous()) {
if (alpha) {
mtl_setBytes(computeEncoder, getMPSScalar(*alpha, iter.common_dtype()), 3);
mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), 3);
}
if (cast_needed) {
std::array<int, 4> size_and_types = {static_cast<int>(c10::elementSize(input.scalar_type())),
Expand All @@ -1117,7 +1124,7 @@ static dispatch_data_t getSectionData(const std::string& name) {
static_cast<int>(out.scalar_type())};
if (alpha) {
mtl_setArgs<3>(computeEncoder,
getMPSScalar(*alpha, iter.common_dtype()),
getMPSScalar(*alpha, alpha_type),
iter.shape(),
iter.strides(0),
iter.strides(1),
Expand Down
31 changes: 31 additions & 0 deletions aten/src/ATen/native/mps/kernels/ActivationKernel.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
#include <c10/metal/indexing.h>
#include <c10/metal/special_math.h>
#include <metal_stdlib>
using namespace metal;
using namespace c10::metal;

struct hardshrink_functor {
template <typename T>
inline T operator()(const T x, const T lambda) {
return (x >= -lambda && x <= lambda) ? T(0) : x;
}
};

struct hardshrink_backward_functor {
template <typename T>
inline T operator()(const T grad_output, const T x, const T lambda) {
return (x >= -lambda && x <= lambda) ? T(0) : grad_output;
}
};

REGISTER_UNARY_ALPHA_OP(hardshrink, float, float, float);
REGISTER_UNARY_ALPHA_OP(hardshrink, half, half, half);
#if __METAL_VERSION__ >= 310
REGISTER_UNARY_ALPHA_OP(hardshrink, bfloat, bfloat, bfloat);
#endif

REGISTER_BINARY_ALPHA_OP(hardshrink_backward, float, float, float);
REGISTER_BINARY_ALPHA_OP(hardshrink_backward, half, half, half);
#if __METAL_VERSION__ >= 310
REGISTER_BINARY_ALPHA_OP(hardshrink_backward, bfloat, bfloat, bfloat);
#endif
66 changes: 33 additions & 33 deletions aten/src/ATen/native/mps/kernels/BinaryKernel.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,35 +362,35 @@ REGISTER_OPMATH_FLOAT_BINARY_OP(remainder);
REGISTER_INTEGER_BINARY_OP(remainder);
REGISTER_OPMATH_FLOAT_BINARY_OP(fmod);
REGISTER_INTEGER_BINARY_OP(fmod);
REGISTER_BINARY_ALPHA_OP(add_alpha, long, long);
REGISTER_BINARY_ALPHA_OP(add_alpha, int, int);
REGISTER_BINARY_ALPHA_OP(add_alpha, float, float);
REGISTER_BINARY_ALPHA_OP(add_alpha, half, half);
REGISTER_BINARY_ALPHA_OP(add_alpha, short, short);
REGISTER_BINARY_ALPHA_OP(add_alpha, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(add_alpha, char, char);
REGISTER_BINARY_ALPHA_OP(add_alpha, bool, bool);
REGISTER_BINARY_ALPHA_OP(sub_alpha, long, long);
REGISTER_BINARY_ALPHA_OP(sub_alpha, int, int);
REGISTER_BINARY_ALPHA_OP(sub_alpha, float, float);
REGISTER_BINARY_ALPHA_OP(sub_alpha, half, half);
REGISTER_BINARY_ALPHA_OP(sub_alpha, short, short);
REGISTER_BINARY_ALPHA_OP(sub_alpha, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(sub_alpha, char, char);
REGISTER_BINARY_ALPHA_OP(sub_alpha, bool, bool);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, long, long);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, int, int);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, float, float);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, half, half);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, short, short);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, char, char);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, bool, bool);
REGISTER_BINARY_ALPHA_OP(add_alpha, long, long, long);
REGISTER_BINARY_ALPHA_OP(add_alpha, int, int, int);
REGISTER_BINARY_ALPHA_OP(add_alpha, float, float, float);
REGISTER_BINARY_ALPHA_OP(add_alpha, half, half, half);
REGISTER_BINARY_ALPHA_OP(add_alpha, short, short, short);
REGISTER_BINARY_ALPHA_OP(add_alpha, uchar, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(add_alpha, char, char, char);
REGISTER_BINARY_ALPHA_OP(add_alpha, bool, bool, bool);
REGISTER_BINARY_ALPHA_OP(sub_alpha, long, long, long);
REGISTER_BINARY_ALPHA_OP(sub_alpha, int, int, int);
REGISTER_BINARY_ALPHA_OP(sub_alpha, float, float, float);
REGISTER_BINARY_ALPHA_OP(sub_alpha, half, half, half);
REGISTER_BINARY_ALPHA_OP(sub_alpha, short, short, short);
REGISTER_BINARY_ALPHA_OP(sub_alpha, uchar, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(sub_alpha, char, char, char);
REGISTER_BINARY_ALPHA_OP(sub_alpha, bool, bool, bool);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, long, long, long);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, int, int, int);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, float, float, float);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, half, half, half);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, short, short, short);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, uchar, uchar, uchar);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, char, char, char);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, bool, bool, bool);

#if __METAL_VERSION__ >= 310
REGISTER_BINARY_ALPHA_OP(add_alpha, bfloat, bfloat);
REGISTER_BINARY_ALPHA_OP(sub_alpha, bfloat, bfloat);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, bfloat, bfloat);
REGISTER_BINARY_ALPHA_OP(add_alpha, bfloat, bfloat, bfloat);
REGISTER_BINARY_ALPHA_OP(sub_alpha, bfloat, bfloat, bfloat);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, bfloat, bfloat, bfloat);
#endif

// Complex binary functions
Expand All @@ -406,9 +406,9 @@ REGISTER_BINARY_OP(add, float2, float2);
REGISTER_BINARY_OP(add, half2, half2);
REGISTER_BINARY_OP(sub, float2, float2);
REGISTER_BINARY_OP(sub, half2, half2);
REGISTER_BINARY_ALPHA_OP(add_alpha, float2, float2);
REGISTER_BINARY_ALPHA_OP(add_alpha, half2, half2);
REGISTER_BINARY_ALPHA_OP(sub_alpha, float2, float2);
REGISTER_BINARY_ALPHA_OP(sub_alpha, half2, half2);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, float2, float2);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, half2, half2);
REGISTER_BINARY_ALPHA_OP(add_alpha, float2, float2, float2);
REGISTER_BINARY_ALPHA_OP(add_alpha, half2, half2, half2);
REGISTER_BINARY_ALPHA_OP(sub_alpha, float2, float2, float2);
REGISTER_BINARY_ALPHA_OP(sub_alpha, half2, half2, half2);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, float2, float2, float2);
REGISTER_BINARY_ALPHA_OP(lerp_alpha, half2, half2, half2);
38 changes: 19 additions & 19 deletions aten/src/ATen/native/mps/kernels/UnaryKernel.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -402,25 +402,25 @@ kernel void round_decimals_strided(
rint(exp10(float(ndigits)) * input[input_offs]) * exp10(float(-ndigits)));
}

#define INSTANTIATE_ROUND_DECIMALS(DTYPE) \
template \
[[host_name("round_decimals_dense_" #DTYPE "_" #DTYPE)]] kernel void \
round_decimals_dense( \
device DTYPE* output [[buffer(0)]], \
constant DTYPE* input [[buffer(1)]], \
constant long& ndigits [[buffer(2)]], \
uint index [[thread_position_in_grid]]); \
template \
[[host_name("round_decimals_strided_" #DTYPE "_" #DTYPE)]] kernel void \
round_decimals_strided( \
device DTYPE* output [[buffer(0)]], \
constant DTYPE* input [[buffer(1)]], \
constant long* sizes, \
constant long* input_strides, \
constant long* output_strides, \
constant uint& ndim, \
constant long& ndigits [[buffer(6)]], \
uint index)
#define INSTANTIATE_ROUND_DECIMALS(DTYPE) \
template [[host_name("round_decimals_dense_" #DTYPE "_" #DTYPE \
"_long")]] kernel void \
round_decimals_dense( \
device DTYPE* output [[buffer(0)]], \
constant DTYPE* input [[buffer(1)]], \
constant long& ndigits [[buffer(2)]], \
uint index [[thread_position_in_grid]]); \
template [[host_name("round_decimals_strided_" #DTYPE "_" #DTYPE \
"_long")]] kernel void \
round_decimals_strided( \
device DTYPE* output [[buffer(0)]], \
constant DTYPE* input [[buffer(1)]], \
constant long* sizes, \
constant long* input_strides, \
constant long* output_strides, \
constant uint& ndim, \
constant long& ndigits [[buffer(6)]], \
uint index)

INSTANTIATE_ROUND_DECIMALS(float);
INSTANTIATE_ROUND_DECIMALS(half);
Expand Down
27 changes: 27 additions & 0 deletions aten/src/ATen/native/mps/operations/ActivationKernel.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/TensorIterator.h>
#include <ATen/mps/MPSProfiler.h>
#include <ATen/native/Activation.h>
#include <ATen/native/mps/OperationUtils.h>
#include <fmt/format.h>

namespace at::native {

#ifndef PYTORCH_JIT_COMPILE_SHADERS
static auto& lib = mps::MetalShaderLibrary::getBundledLibrary();
#else
#include <ATen/native/mps/ActivationKernel_metallib.h>
#endif

static void hardshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
lib.exec_unary_kernel(iter, "hardshrink", lambda);
}

static void hardshrink_backward_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
lib.exec_binary_kernel(iter, "hardshrink_backward", lambda);
}

REGISTER_DISPATCH(hardshrink_stub, hardshrink_kernel);
REGISTER_DISPATCH(shrink_backward_stub, hardshrink_backward_kernel);

} // namespace at::native
2 changes: 1 addition & 1 deletion aten/src/ATen/native/mps/operations/UnaryKernel.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
REGISTER_DISPATCH(NAME##_stub, NAME##_kernel_mps)

static void round_decimals_kernel(TensorIteratorBase& iter, int64_t decimals) {
lib.exec_unary_kernel(iter, "round_decimals", decimals);
lib.exec_unary_kernel(iter, "round_decimals", Scalar(decimals), ScalarType::Long);
}

REGISTER_UNARY_TI_DISPATCH(exp);
Expand Down
4 changes: 2 additions & 2 deletions aten/src/ATen/native/native_functions.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5172,7 +5172,7 @@
structured_inherits: TensorIteratorBase
device_check: NoCheck # TensorIterator
dispatch:
CPU, CUDA: hardshrink_out
CPU, CUDA, MPS: hardshrink_out

- func: hardshrink(Tensor self, Scalar lambd=0.5) -> Tensor
structured_delegate: hardshrink.out
Expand All @@ -5184,7 +5184,7 @@
structured: True
structured_inherits: TensorIteratorBase
dispatch:
CPU, CUDA: hardshrink_backward_out
CPU, CUDA, MPS: hardshrink_backward_out

- func: hardshrink_backward(Tensor grad_out, Tensor self, Scalar lambd) -> Tensor
structured_delegate: hardshrink_backward.grad_input
Expand Down
Loading
Loading