diff --git a/mlx/backend/cpu/binary_ops.h b/mlx/backend/cpu/binary_ops.h index d50751ce3d..c617c70f92 100644 --- a/mlx/backend/cpu/binary_ops.h +++ b/mlx/backend/cpu/binary_ops.h @@ -33,8 +33,46 @@ DEFAULT_BINARY_OP(LogicalOr, operator||) DEFAULT_BINARY_OP(BitwiseAnd, operator&) DEFAULT_BINARY_OP(BitwiseOr, operator|) DEFAULT_BINARY_OP(BitwiseXor, operator^) -DEFAULT_BINARY_OP(LeftShift, operator<<) -DEFAULT_BINARY_OP(RightShift, operator>>) +// A shift amount that is negative or at least the operand width is undefined +// in C++, so the scalar and vectorized paths disagreed about it and the answer +// depended on the array length. Clamp the amount instead, which gives what +// numpy and pytorch return: zero for a left shift, and the sign bit for a +// right shift. +template +Simd shift_in_range(Simd y) { + auto in_range = y < Simd(sizeof(T) * 8); + if constexpr (std::is_signed_v) { + in_range = in_range && y >= Simd(0); + } + return in_range; +} + +struct LeftShift { + template + Simd operator()(Simd x, Simd y) { + Simd amount = y & Simd(sizeof(T) * 8 - 1); + Simd shifted = x << amount; + return select(shift_in_range(y), shifted, Simd(0)); + } + BINARY_SINGLE() +}; + +struct RightShift { + template + Simd operator()(Simd x, Simd y) { + auto in_range = shift_in_range(y); + Simd amount = select(in_range, y, Simd(sizeof(T) * 8 - 1)); + Simd shifted = x >> amount; + if constexpr (std::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return select(in_range, shifted, Simd(0)); + } + } + BINARY_SINGLE() +}; + DEFAULT_BINARY_OP(Remainder, remainder) DEFAULT_BINARY_OP(Maximum, maximum) DEFAULT_BINARY_OP(Minimum, minimum) diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh index b0b7962807..75111b0eba 100644 --- a/mlx/backend/cuda/device/binary_ops.cuh +++ b/mlx/backend/cuda/device/binary_ops.cuh @@ -269,17 +269,38 @@ struct BitwiseXor { }; }; +// A shift amount that is negative or at least the operand width is undefined, +// so clamp it to what numpy and pytorch return: zero for a left shift, and the +// sign bit for a right shift. +template +__device__ bool shift_in_range(T y) { + bool in_range = y < static_cast(sizeof(T) * 8); + if constexpr (cuda::std::is_signed_v) { + in_range = in_range && y >= 0; + } + return in_range; +} + struct LeftShift { template __device__ T operator()(T x, T y) { - return x << y; + T amount = y & static_cast(sizeof(T) * 8 - 1); + return shift_in_range(y) ? static_cast(x << amount) : static_cast(0); }; }; struct RightShift { template __device__ T operator()(T x, T y) { - return x >> y; + bool in_range = shift_in_range(y); + T amount = in_range ? y : static_cast(sizeof(T) * 8 - 1); + T shifted = x >> amount; + if constexpr (cuda::std::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return in_range ? shifted : static_cast(0); + } }; }; diff --git a/mlx/backend/metal/kernels/binary_ops.h b/mlx/backend/metal/kernels/binary_ops.h index 863d6369e2..0bcee104ce 100644 --- a/mlx/backend/metal/kernels/binary_ops.h +++ b/mlx/backend/metal/kernels/binary_ops.h @@ -303,17 +303,38 @@ struct BitwiseXor { }; }; +// A shift amount that is negative or at least the operand width is undefined, +// so clamp it to what numpy and pytorch return: zero for a left shift, and the +// sign bit for a right shift. +template +inline bool shift_in_range(T y) { + bool in_range = y < static_cast(sizeof(T) * 8); + if constexpr (metal::is_signed_v) { + in_range = in_range && y >= 0; + } + return in_range; +} + struct LeftShift { template T operator()(T x, T y) thread { - return x << y; + T amount = y & static_cast(sizeof(T) * 8 - 1); + return shift_in_range(y) ? static_cast(x << amount) : static_cast(0); }; }; struct RightShift { template T operator()(T x, T y) thread { - return x >> y; + bool in_range = shift_in_range(y); + T amount = in_range ? y : static_cast(sizeof(T) * 8 - 1); + T shifted = x >> amount; + if constexpr (metal::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return in_range ? shifted : static_cast(0); + } }; }; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index dd0a93001a..ac308a0ef7 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3448,6 +3448,34 @@ def test_bitwise_ops(self): out_np = getattr(np, op)(a_np, b_np) self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) + # A shift amount that is negative or at least the operand width is + # undefined in C++, so the scalar and vectorized paths disagreed and + # the answer depended on the length of the array. + for t in types: + nt = np.array(mx.zeros(1, dtype=t)).dtype + info = np.iinfo(nt) + width = t.size * 8 + amounts = [0, 1, width - 1, width, width + 3] + values = [1, 16, info.max] + if info.min < 0: + amounts += [-1, -width] + values += [-16, info.min] + for n in (1, 7, 8, 100): + for v in values: + for sh in amounts: + a_np = np.full(n, v, nt) + b_np = np.full(n, sh, nt) + a_mlx = mx.array(a_np) + b_mlx = mx.array(b_np) + for op in ("left_shift", "right_shift"): + self.assertTrue( + np.array_equal( + np.array(getattr(mx, op)(a_mlx, b_mlx)), + getattr(np, op)(a_np, b_np), + ), + msg=f"{op} {t} n={n} {v} by {sh}", + ) + for t in types: a_mlx = a.astype(t) a_np = np.array(a_mlx)