fix(conv): use the fp32 fallback when set_precision requests float32 or float64

aether/layers/conv.py

@@ -29,18 +29,21 @@ class Conv2d(Layer):
     Warns
     -----
     UserWarning
-        If the configured ``precision_policy`` requests a compute precision
-        that does not match the active execution path (e.g., requesting a dtype
-        other than ``float16`` when bound to the WMMA matrix-core path, or
-        requesting non-``float32`` when falling back to strided CPU/GPU kernels).
+        If the configured ``precision_policy`` requests ``float32`` or
+        ``float64`` while the WMMA matrix-core path is available: the layer
+        then binds the much slower strided float32 fallback (``float64`` is not
+        honored, float32 is the most it computes in).
+        If it requests a dtype the bound path cannot deliver otherwise (e.g.
+        ``float16`` or ``float64`` on the strided CPU/GPU fallback).
 
     Notes
     -----
     Expected input shape: ``(batch_size, height, width, in_channels)``
     Output shape: ``(batch_size, out_height, out_width, out_channels)``
 
-    Conv2d's compute dtype is bounded to fp16 as WMMA kernels can't take >fp16
-    calculations, instead layer will use the fallback vectorized path.
+    Unless a precision is set, on a GPU the layer uses fp16 WMMA kernels
+    (float32 accumulate). WMMA can't take >fp16 operands, so an explicit
+    ``float32`` request uses the slower fallback vectorized path instead.
     """
 
     def __init__(self, in_channels, out_channels = 1, filter_size = (3, 3), stride = (1, 1), padding = "same", l1=(), l2=()):
@@ -108,7 +111,8 @@ def build(self, input_shape: tuple[int, ...], seed: int | None = None) -> tuple[
         return self.output_shape
     def _compile_for_device(self, device):
         """Triggered by Model.to(device) to bind the matrix-core path when available."""
-        if device == 'cupy' and conv_kernel.get_is_conv_gpu_available():
+        # An explicit float32/float64 request opts out of the fp16 matrix cores.
+        if self._wmma_available(device) and self._requested_dtype() not in ('float32', 'float64'):
             self.forward = self._forward_gpu
             self.backward = self._backward_gpu
         else:
@@ -120,17 +124,24 @@ def _compile_for_device(self, device):
         # Covers set_precision() -> to(); the binding is only knowable here.
         self._warn_if_precision_ignored()
 
+    @staticmethod
+    def _wmma_available(device):
+        return device == 'cupy' and conv_kernel.get_is_conv_gpu_available()
+
+    def _requested_dtype(self):
+        return self.precision_policy.compute_dtype_name if self.precision_policy else None
+
     def _apply_precision(self, policy):
         """
-        Called on Model.set_precision(). Conv2d's kernels pin their own operand
-        dtype, so the policy is stored for introspection/serialization only --
-        _warn_if_precision_ignored is what keeps that from being silent.
+        Called on Model.set_precision(). The policy is stored for serialization,
+        and a compiled layer re-binds its path, since an explicit float32/float64
+        request moves it off the matrix cores (and float16 / None back on).
         """
         self.precision_policy = policy or config.DTypePolicy()
         # Before _compile_for_device runs there is no binding to check against,
         # so defer to it rather than warning about a placeholder.
         if self._compiled_device is not None:
-            self._warn_if_precision_ignored()
+            self._compile_for_device(self._compiled_device)
 
     def _effective_compute_dtype(self):
         """The dtype this layer actually multiplies in, read off the live binding."""
@@ -142,17 +153,19 @@ def _effective_compute_dtype(self):
     def _warn_if_precision_ignored(self):
         """
         Warn once per (bound path, requested dtype) when the precision policy
-        asks for something this layer's kernels can't deliver. Fires at bind
+        asks for something this layer can't deliver as requested. Fires at bind
         time only -- never from forward/backward.
         """
-        requested = (
-            self.precision_policy.compute_dtype_name if self.precision_policy else None
-        )
+        requested = self._requested_dtype()
         if requested is None:
             return
 
         effective = self._effective_compute_dtype()
-        if requested == effective:
+        # float32/float64 deliberately left the matrix cores: honored, but slow.
+        slow_by_request = (
+            requested in ('float32', 'float64') and self._wmma_available(self._compiled_device)
+        )
+        if requested == effective and not slow_by_request:
             self._precision_warned_key = None
             return
 
@@ -163,6 +176,20 @@ def _warn_if_precision_ignored(self):
             return
         self._precision_warned_key = key
 
+        if slow_by_request:
+            cap = (
+                " Conv2d computes in float32 at most, so float64 is not honored."
+                if requested == 'float64' else ""
+            )
+            warnings.warn(
+                f"[aether] Conv2d: precision policy requests {requested!r}, so this "
+                f"layer uses the strided float32 fallback path instead of the "
+                f"float16 matrix-core path.{cap} THIS IS MUCH SLOWER; for the fast "
+                f"path use set_precision('float16') or leave it unset.",
+                UserWarning,
+            )
+            return
+
         path = (
             "the matrix-core path, whose WMMA kernels take float16 operands at "
             "the hardware level (float32 accumulate, float32 output)"

aether/model.py

@@ -440,7 +440,10 @@ def set_precision(self, compute_dtype):
 
         Args:
             compute_dtype (str): The target floating-point precision -- one of
-                'float16', 'float32', or 'float64'.
+                'float16', 'float32', or 'float64'. If never set, layers use
+                mixed precision where their kernels allow it; 'float32' asks
+                for strict fp32, so Conv2d uses its slower fallback on GPU
+                (and warns).
         Raises:
             TypeError: If `compute_dtype` is not a str or None.
             ValueError: If `compute_dtype` is not one of the currently supported precisions.

tests/layers/conv2d/test_precision.py

@@ -2,6 +2,7 @@
 
 import aether.config as config
 import tests.base_case as base_case
+from aether.custom_kernels import conv_kernel
 from aether.layers.conv import Conv2d
 from tests.layers.conv2d.base import ConvLayerTestBase
 
@@ -35,8 +36,9 @@ def test_policy_matching_the_bound_path_is_silent(self):
 
     def test_mismatched_policy_warns_once_naming_both_dtypes(self):
         """float64 is never an effective dtype, so it mismatches on either path."""
-        effective = self.layer._effective_compute_dtype()
         caught = self._capture(self._apply(self.layer, "float64"))
+        # After applying: on a WMMA GPU float64 moves the layer to the fallback.
+        effective = self.layer._effective_compute_dtype()
 
         self.assertEqual(len(caught), 1)
         self.assertIn("float64", caught[0])
@@ -77,5 +79,68 @@ def test_layer_still_computes_under_an_ignored_policy(self):
         self.assertEqual(out.shape, (self.CONV_SHAPE[0], 28, 28, self.OUT_CHANNELS))
         self.assertTrue(bool(self.xp.all(self.xp.isfinite(out))))
 
+    def _require_wmma(self):
+        if self.backend_name != "cupy" or not conv_kernel.get_is_conv_gpu_available():
+            self.skipTest("needs a GPU with the WMMA conv path")
+
+    def _is_wmma(self, layer):
+        return layer._effective_compute_dtype() == "float16"
+
+    def test_explicit_float32_binds_fallback_with_slow_warning(self):
+        """float32 opts out of WMMA, in either order of set_precision() and compile."""
+        self._require_wmma()
+        # compile -> set_precision
+        self.assertTrue(self._is_wmma(self.layer))
+        caught = self._capture(self._apply(self.layer, "float32"))
+        self.assertFalse(self._is_wmma(self.layer))
+        self.assertEqual(len(caught), 1)
+        self.assertTrue(caught[0].startswith("[aether]"))
+        self.assertIn("MUCH SLOWER", caught[0])
+
+        # set_precision -> compile
+        layer = self._make_uncompiled_conv()
+        self.assertEqual(self._capture(self._apply(layer, "float32")), [])
+        caught = self._capture(lambda: [layer._compile_for_device("cupy") for _ in range(2)])
+        self.assertFalse(self._is_wmma(layer))
+        self.assertEqual(len(caught), 1)
+        self.assertIn("MUCH SLOWER", caught[0])
+
+    def test_explicit_float64_binds_fallback_and_says_float32(self):
+        self._require_wmma()
+        caught = self._capture(self._apply(self.layer, "float64"))
+        self.assertFalse(self._is_wmma(self.layer))
+        self.assertEqual(len(caught), 1)
+        self.assertIn("MUCH SLOWER", caught[0])
+        self.assertIn("float32 at most", caught[0])
+
+    def test_unset_and_float16_stay_on_wmma_silently(self):
+        self._require_wmma()
+        for requested in (None, "float16"):
+            with self.subTest(requested=requested):
+                layer = self._make_uncompiled_conv()
+                layer._apply_precision(config.DTypePolicy(requested))
+                caught = self._capture(lambda layer=layer: layer._compile_for_device("cupy"))
+                self.assertEqual(caught, [])
+                self.assertTrue(self._is_wmma(layer))
+
+    def test_float32_then_float16_returns_to_wmma(self):
+        self._require_wmma()
+        self._capture(self._apply(self.layer, "float32"))
+        self.assertEqual(self._capture(self._apply(self.layer, "float16")), [])
+        self.assertTrue(self._is_wmma(self.layer))
+
+    def test_float32_forward_is_bit_exact_with_the_fallback(self):
+        """Under float32 the GPU layer really runs the fp32 path, not fp16 WMMA."""
+        self._require_wmma()
+        x = self.test_images.astype(self.xp.float32)
+        wmma_out = self.layer.forward(x, training=False).copy()
+        self._capture(self._apply(self.layer, "float32"))
+
+        out = self.layer.forward(x, training=False).copy()
+        expected = self.layer._forward_fallback(x, training=False)
+        self.assertTrue(bool(self.xp.array_equal(out, expected)))
+        # Guard: fp16 WMMA rounds differently, so the check can tell the paths apart.
+        self.assertFalse(bool(self.xp.array_equal(wmma_out, expected)))
+
 
 base_case.register_test_suites(globals(), TestConvPrecision)