Add MLX_CONV_WINOGRAD switch to disable winograd for conv2d

docs/src/usage/environment_variables.rst

@@ -102,6 +102,12 @@ Metal
    Enable the faster Metal CPU/GPU synchronization path. The default is ``0``.
    This requires Metal 3.2 or later (macOS 15 or later, or iOS 18 or later).
 
+.. envvar:: MLX_CONV_WINOGRAD
+
+   Allow the Winograd algorithm for 3x3 ``conv2d``. The default is ``1``. Set
+   it to ``0`` to use the GEMM convolution for all inputs. See
+   :doc:`precision`.
+
 Advanced tuning
 ---------------
 

docs/src/usage/precision.rst

@@ -19,3 +19,15 @@ To keep these operations in full ``float32``, set
 
 Which operations take the reduced-precision path, and how large the
 difference is, depends on the backend and the hardware.
+
+On Metal, ``conv2d`` can use the Winograd algorithm for 3x3 convolutions
+with stride 1. Winograd is faster for large inputs, but its rounding error
+in ``float32`` can be approximately 10 times larger. MLX selects Winograd
+from the total input size, which includes the batch size. Thus, the same
+image can give different results in a batch of 3 and in a batch of 1.
+
+To disable Winograd, set :envvar:`MLX_CONV_WINOGRAD` to ``0``:
+
+.. code-block:: shell
+
+  MLX_CONV_WINOGRAD=0 python my_script.py

mlx/backend/metal/conv.cpp

@@ -1450,7 +1450,7 @@ void dispatch_conv_2D_gpu(
   if (!conv_params.flip && is_stride_one && is_kdil_one && is_idil_one &&
       conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
       conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large &&
-      channels_large) {
+      channels_large && env::get_var("MLX_CONV_WINOGRAD", 1)) {
     // Only use winograd conv when having enough memory.
     if (int n_step = winograd_batch_step(d, in, conv_params); n_step > 0) {
       return winograd_conv_2D_gpu(

python/tests/test_conv.py

@@ -1481,6 +1481,25 @@ def run(x, w, env={}):
                 self.assertFalse(np.array_equal(untiled, capped))
                 self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3))
 
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_conv2d_winograd_disable(self):
+        # N * iH * iW >= 4096 selects Winograd only from a batch of 3.
+        np.random.seed(0)
+        x = mx.array(np.random.normal(size=(1, 40, 40, 128)).astype(np.float32))
+        w = mx.array(
+            (np.random.normal(size=(128, 3, 3, 128)) * 0.05).astype(np.float32)
+        )
+        xb = mx.repeat(x, 3, axis=0)
+
+        single = np.array(mx.conv2d(x, w, padding=1))
+        winograd = np.array(mx.conv2d(xb, w, padding=1))
+        self.assertFalse(np.array_equal(winograd[0], single[0]))
+
+        with mlx_tests.scoped_env(MLX_CONV_WINOGRAD="0"):
+            gemm = np.array(mx.conv2d(xb, w, padding=1))
+        for i in range(3):
+            self.assertTrue(np.array_equal(gemm[i], single[0]))
+
     def test_conv2d_large_filter_small_channels(self):
         x = mx.random.normal(shape=(1, 181, 181, 1))
         w = mx.random.normal(shape=(1, 182, 182, 1))