diff --git a/mlx/backend/metal/kernels/copy.h b/mlx/backend/metal/kernels/copy.h index cf22347ee5..95ed69b760 100644 --- a/mlx/backend/metal/kernels/copy.h +++ b/mlx/backend/metal/kernels/copy.h @@ -9,11 +9,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } } @@ -27,11 +27,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } } @@ -46,11 +46,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } } @@ -65,11 +65,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } } @@ -81,7 +81,7 @@ template constant const int64_t& src_stride [[buffer(3)]], uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); - dst[index] = static_cast(src[src_idx]); + dst[index] = cast_to(src[src_idx]); } template @@ -93,7 +93,7 @@ template uint2 grid_dim [[threads_per_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * index.y; - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -106,7 +106,7 @@ template auto src_idx = elem_to_loc_3(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -123,14 +123,14 @@ template if (N == 1) { IdxT dst_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); return; } auto xshape = src_shape[ndim - 1]; IdxT dst_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z); auto src_xstride = src_strides[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[dst_idx + i] = static_cast(src[src_idx]); + dst[dst_idx + i] = cast_to(src[src_idx]); src_idx += src_xstride; } } @@ -144,7 +144,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -156,7 +156,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -168,7 +168,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -187,14 +187,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } @@ -211,7 +211,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -225,7 +225,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -239,7 +239,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -262,14 +262,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } diff --git a/mlx/backend/metal/kernels/reduction/reduce_all.h b/mlx/backend/metal/kernels/reduction/reduce_all.h index e0d08392c0..47ad63fbbd 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_all.h +++ b/mlx/backend/metal/kernels/reduction/reduce_all.h @@ -37,13 +37,13 @@ template < for (IdxT b = 0; b < blocks; b++) { for (int i = 0; i < N_READS; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } in += lsize.x * N_READS; } if (extra > 0) { for (int i = 0; i < extra; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } } diff --git a/mlx/backend/metal/kernels/reduction/reduce_col.h b/mlx/backend/metal/kernels/reduction/reduce_col.h index c109faf0bc..b1546adb55 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_col.h +++ b/mlx/backend/metal/kernels/reduction/reduce_col.h @@ -43,13 +43,13 @@ template row = in + loop.location(); if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -125,7 +125,7 @@ template for (IdxT r = gid.z * lsize.y + lid.y; r < total_rows; r += lsize.y * gsize.z) { row = in + loop.location(); - total = op(static_cast(*row), total); + total = op(cast_to(*row), total); loop.next(lsize.y * gsize.z, reduce_shape, reduce_strides); } @@ -207,13 +207,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -352,13 +352,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); diff --git a/mlx/backend/metal/kernels/reduction/reduce_row.h b/mlx/backend/metal/kernels/reduction/reduce_row.h index 936d75bb52..90b2226460 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_row.h +++ b/mlx/backend/metal/kernels/reduction/reduce_row.h @@ -34,7 +34,7 @@ METAL_FUNC void per_thread_row_reduce( for (int i = 0; i < blocks; i++) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } inputs[j] += lsize_x * N_READS; @@ -46,13 +46,13 @@ METAL_FUNC void per_thread_row_reduce( if (index + N_READS <= extra) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } else { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; index + i < extra; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } diff --git a/mlx/backend/metal/kernels/utils.h b/mlx/backend/metal/kernels/utils.h index 266f27e91c..f15928282f 100644 --- a/mlx/backend/metal/kernels/utils.h +++ b/mlx/backend/metal/kernels/utils.h @@ -446,3 +446,27 @@ template struct ConditionalType { using type = T; }; + +/////////////////////////////////////////////////////////////////////////////// +// Type casting utils +/////////////////////////////////////////////////////////////////////////////// + +template +inline U cast_to(T val) { + return static_cast(val); +} + +template <> +inline bool cast_to(float val) { + return (as_type(val) & 0x7FFFFFFF) != 0; +} + +template <> +inline bool cast_to(bfloat16_t val) { + return (as_type(val) & 0x7FFF) != 0; +} + +template <> +inline bool cast_to(complex64_t val) { + return cast_to(val.real) || cast_to(val.imag); +} diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d569dbd3c7..60c0052532 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1299,6 +1299,21 @@ def test_any(self): self.assertEqual(mx.any(a, axis=0).tolist(), [True, False]) self.assertEqual(mx.any(a, axis=1).tolist(), [True, False]) + def test_subnormal_bool_cast(self): + f32_sub = mx.array(np.array([0x00000001], dtype=np.uint32)).view(mx.float32) + f16_sub = mx.array(np.array([0x0001], dtype=np.uint16)).view(mx.float16) + bf16_sub = mx.array(np.array([0x0001], dtype=np.uint16)).view(mx.bfloat16) + + self.assertTrue(f32_sub.astype(mx.bool_).item()) + self.assertTrue(f16_sub.astype(mx.bool_).item()) + self.assertTrue(bf16_sub.astype(mx.bool_).item()) + self.assertTrue(mx.any(f32_sub).item()) + self.assertTrue(mx.any(f16_sub).item()) + self.assertTrue(mx.any(bf16_sub).item()) + self.assertTrue(mx.all(f32_sub).item()) + self.assertTrue(mx.all(f16_sub).item()) + self.assertTrue(mx.all(bf16_sub).item()) + def test_stop_gradient(self): def func(x): return mx.sum(2 * x + mx.stop_gradient(3 * x))