Repository: jax-ml/jax
Pull request #24510 (2024-10-24 17:47:04)

Allow dot algorithms in default_matmul_precision config
I had already implemented handling this config value when it was set to an algorithm name, but these values weren't allowed as input to this config. This was an oversight on my part - whoops!

This doesn't support passing arbitrary algorithms using the config, but all the currently known presets are now allowed.

cc @sbodenstein 


Patch:
diff --git a/jax/_src/config.py b/jax/_src/config.py
index a05e6e190..808048324 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1469,7 +1469,16 @@ numpy_rank_promotion = enum_state(
 
 default_matmul_precision = optional_enum_state(
     name='jax_default_matmul_precision',
-    enum_values=['default', 'high', 'highest', 'bfloat16', 'tensorfloat32', 'float32'],
+    enum_values=[
+        # Legacy precision API values
+        'default', 'high', 'highest', 'bfloat16', 'tensorfloat32', 'float32',
+        # Dot algorithm presets
+        'ANY_F8_ANY_F8_F32', 'ANY_F8_ANY_F8_F32_FAST_ACCUM', 'ANY_F8_ANY_F8_ANY',
+        'ANY_F8_ANY_F8_ANY_FAST_ACCUM', 'F16_F16_F16', 'F16_F16_F32',
+        'BF16_BF16_BF16', 'BF16_BF16_F32', 'BF16_BF16_F32_X3',
+        'BF16_BF16_F32_X6', 'TF32_TF32_F32', 'TF32_TF32_F32_X3', 'F32_F32_F32',
+        'F64_F64_F64',
+    ],
     default=None,
     help=('Control the default matmul and conv precision for 32bit inputs.\n\n'
 
@@ -1486,7 +1495,12 @@ default_matmul_precision = optional_enum_state(
           'convolution on 32bit inputs. The levels roughly describe the '
           "precision at which scalar products are computed. The 'bfloat16' "
           "option is the fastest and least precise; 'float32' is similar to "
-          "full float32 precision; 'tensorfloat32' is intermediate.\n\n"),
+          "full float32 precision; 'tensorfloat32' is intermediate.\n\n"
+
+          'This parameter can also be used to specify an accumulation '
+          '"algorithm" for functions that perform matrix multiplications, like '
+          ':func:`jax.lax.dot`. To specify an algorithm, set this option to '
+          'the name of a :class:`~jax.lax.DotAlgorithmPreset`.\n\n'),
     update_global_hook=lambda val: \
       _update_global_jit_state(default_matmul_precision=val),
     update_thread_local_hook=lambda val: \


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 4d20240e1..79ad9fcfa 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1156,6 +1156,18 @@ class LaxTest(jtu.JaxTestCase):
     lhs, rhs = rng(lhs_shape, np.float16), rng(rhs_shape, np.float16)
     self.assertEqual(fun(lhs, rhs).dtype, np.float16)
 
+  def testDotAlgorithmConfig(self):
+    lhs_shape = (3, 4)
+    rhs_shape = (4, 3)
+    rng = jtu.rand_default(self.rng())
+    lhs, rhs = rng(lhs_shape, np.float32), rng(rhs_shape, np.float32)
+
+    expected = ("algorithm = <lhs_precision_type = f32, rhs_precision_type = "
+                "f32, accumulation_type = f32")
+    with jax.default_matmul_precision("F32_F32_F32"):
+      hlo = jax.jit(lax.dot).lower(lhs, rhs).as_text()
+      self.assertRegex(hlo, expected)
+
   @jtu.sample_product(
     [dict(lhs_shape=lhs_shape, rhs_shape=rhs_shape)
      for lhs_shape in [(3,), (4, 3)] for rhs_shape in [(3,), (3, 6)]],