diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 926f31f05a..08ece0f10a 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -948,7 +948,13 @@ void winograd_conv_2D_gpu( /*b_cols = */ conv_params.O, /*a_transposed = */ false, /*b_transposed = */ false, - /*copies = */ empty_copies); + /*copies = */ empty_copies, + /*batch_shape = */ {}, + /*A_batch_stride = */ {}, + /*B_batch_stride = */ {}, + // The winograd transforms amplify whatever the gemm rounds away, so + // this one stays at full float32 even when tf32 is allowed elsewhere. + /*allow_tf32 = */ false); } // Do output transform diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 8c0b46d6de..d3600781fc 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -862,7 +862,8 @@ void steel_matmul_axpby( Strides B_batch_stride /* = {} */, Strides C_batch_stride /* = {} */, float alpha /* = 1.0f */, - float beta /* = 0.0f */) { + float beta /* = 0.0f */, + bool allow_tf32 /* = true */) { if (batch_shape.empty()) { ///////////////////////////////////////////////////////////////////////////// // Check and collapse batch dimensions @@ -916,7 +917,7 @@ void steel_matmul_axpby( int64_t matrix_size = static_cast(M) * N; bool use_nax = metal::is_nax_available() && !issubdtype(a.dtype(), complexfloating) && - (env::enable_tf32() || a.dtype() != float32); + ((allow_tf32 && env::enable_tf32()) || a.dtype() != float32); char devc = d.get_architecture().back(); int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024; diff --git a/mlx/backend/metal/matmul.h b/mlx/backend/metal/matmul.h index 218664b1ff..fd1f0a65b8 100644 --- a/mlx/backend/metal/matmul.h +++ b/mlx/backend/metal/matmul.h @@ -100,7 +100,8 @@ void steel_matmul_axpby( Strides B_batch_stride = {}, Strides C_batch_stride = {}, float alpha = 1.0f, - float beta = 0.0f); + float beta = 0.0f, + bool allow_tf32 = true); inline void steel_matmul( const Stream& s, @@ -119,7 +120,8 @@ inline void steel_matmul( std::vector& copies, Shape batch_shape = {}, Strides A_batch_stride = {}, - Strides B_batch_stride = {}) { + Strides B_batch_stride = {}, + bool allow_tf32 = true) { return steel_matmul_axpby( /* const Stream& s = */ s, /* metal::Device& d = */ d, @@ -138,7 +140,11 @@ inline void steel_matmul( /* std::vector& copies = */ copies, /* Shape batch_shape = */ batch_shape, /* Strides A_batch_stride = */ A_batch_stride, - /* Strides B_batch_stride = */ B_batch_stride); + /* Strides B_batch_stride = */ B_batch_stride, + /* Strides C_batch_stride = */ {}, + /* float alpha = */ 1.0f, + /* float beta = */ 0.0f, + /* bool allow_tf32 = */ allow_tf32); } } // namespace mlx::core diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index c5f9a2c1b2..054f98fa19 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -2,6 +2,9 @@ import math import os +import subprocess +import sys +import textwrap import unittest from itertools import permutations @@ -48,6 +51,34 @@ def test_numpy_conv(self): self.assertEqual(c_mx.shape, c_np.shape) self.assertTrue(np.allclose(c_mx, c_np, atol=atol)) + def test_conv_2d_winograd_float32_precision(self): + # The aligned conv takes winograd and the unaligned one does not, so + # the two have to agree to float32 precision. + script = textwrap.dedent(""" + import mlx.core as mx + + # Winograd also wants N * iH * iW >= 4096 and C + O >= 256. + N, S, K, C = 4, 32, 3, 128 + x = mx.random.normal((N, S, S, C)) + w = mx.random.normal((C, K, K, C)) + aligned = mx.conv2d(x, w, padding=1) + + # The same conv with an input channel count winograd declines. + xu = mx.pad(x, [(0, 0), (0, 0), (0, 0), (0, 16)]) + wu = mx.pad(w, [(0, 0), (0, 0), (0, 0), (0, 16)]) + unaligned = mx.conv2d(xu, wu, padding=1) + + mx.eval(aligned, unaligned) + print(mx.abs(aligned - unaligned).max().item()) + """) + # tf32 is off for the whole suite and enable_tf32() caches the read. + env = dict(os.environ, MLX_ENABLE_TF32="1") + out = subprocess.check_output( + [sys.executable, "-c", script], env=env, text=True + ) + gap = float(out.strip()) + self.assertLess(gap, 1e-2, f"winograd and gemm disagree by {gap}") + def test_conv_1d_groups_flipped(self): x = mx.broadcast_to(mx.arange(5).astype(mx.float32), (2, 5)).T w = mx.broadcast_to(mx.arange(4).astype(mx.float32), (2, 4))