Repository: jax-ml/jax
Pull request #27679 (2025-04-02 20:40:45)

Support `conv` `unfused_flops` in roofline.
Support `conv` `unfused_flops` in roofline.

Since calculating flops is non-trivial, we don't test all the cases currently tested by `test_conv_general_dilated_unfused_hbm_bytes`. Instead, we test behaviors more directly.



Patch:
diff --git a/jax/experimental/roofline/rooflines.py b/jax/experimental/roofline/rooflines.py
index bc8d65e96..63a2d7cc4 100644
--- a/jax/experimental/roofline/rooflines.py
+++ b/jax/experimental/roofline/rooflines.py
@@ -14,6 +14,7 @@
 from collections import defaultdict
 from dataclasses import replace
 import itertools as it
+from typing import Sequence
 import numpy as np
 
 from jax._src import ad_util
@@ -35,7 +36,7 @@ from jax._src.lax import (
 from jax.experimental import roofline
 from jax.experimental import shard_map
 
-
+# One FMA (Fused Multiply Add) takes 2 flops to compute.
 _FMA_FLOPS_FACTOR = 2
 
 for prim in it.chain(
@@ -179,16 +180,208 @@ def _dot_general_roofline(
       unfused_hbm_bytes=hbm_bytes,
   )
 
+
+def _get_spatial_valid_position_count_for_one_dim(
+    window_dim_stride: int,
+    base_dilation: int,
+    window_dilation: int,
+    kernel_limit: int,
+    input_limit: int,
+    output_limit: int,
+    padding: tuple[int, int],
+) -> int:
+  """Gets the valid position count for conv for a single spatial dimension.
+
+  Args:
+    window_dim_stride: The stride of the window along this dimension.
+    base_dilation: The base dilation factor along this dimension.
+    window_dilation: The window dilation factor along this dimension.
+    kernel_limit: The size of the kernel along this dimension.
+    input_limit: The size of the input along this dimension.
+    output_limit: The size of the output along this dimension.
+    padding: The padding applied to the input along this dimension.
+  """
+  padding_low = padding[0]
+  padding_high = padding[1]
+
+  # These two conditions will create an N^2 iteration pattern with only N
+  # valid elements. This is a performance optimization and produces the same
+  # result as the whole loop.
+  if (
+      input_limit == output_limit
+      and kernel_limit == output_limit
+      and input_limit == base_dilation
+      and window_dilation == 1
+      and max(1, input_limit - 1) == window_dim_stride
+      and padding_low == 0
+      and padding_high == 0
+  ):
+    return input_limit
+
+  if (
+      input_limit == 1
+      and kernel_limit == output_limit
+      and window_dilation == 1
+      and base_dilation == 1
+      and window_dim_stride == 1
+      and padding_low == output_limit - 1
+      and padding_high == output_limit - 1
+  ):
+    return output_limit
+
+  valid_position_count = 0
+  # Loop over each point in the kernel
+  for kernel_idx in range(kernel_limit):
+
+    # Skip loop for trivial stride and base_dilation
+    if window_dim_stride == 1 and base_dilation == 1:
+      undilated_index_base = padding_low - kernel_idx * window_dilation
+      upper_limit = min(
+          input_limit + undilated_index_base,
+          output_limit,
+      )
+      lower_limit = max(0, undilated_index_base)
+
+      valid_position_count += max(upper_limit - lower_limit, 0)
+      continue
+
+    # Loop over each point in the output
+    for output_idx in range(output_limit):
+      # Calculate lhs (input) index without taking base dilation into account
+      undilated_index = (
+          output_idx * window_dim_stride
+          - padding_low
+          + kernel_idx * window_dilation
+      )
+      # Calculate the actual lhs (input) index after dilation
+      lhs_spatial_index = int(undilated_index / base_dilation)
+
+      # Skip if the lhs (input) index is to be dilated.
+      if undilated_index != lhs_spatial_index * base_dilation:
+        continue
+      # Skip if input index is not in bound.
+      if lhs_spatial_index < 0 or lhs_spatial_index >= input_limit:
+        continue
+
+      valid_position_count += 1
+  return valid_position_count
+
+
+def _get_spatial_valid_position_count(
+    dnums: convolution.ConvDimensionNumbers,
+    lhs: roofline.RooflineShape,
+    rhs: roofline.RooflineShape,
+    out: roofline.RooflineShape,
+    window_strides: Sequence[int],
+    padding: Sequence[tuple[int, int]],
+    lhs_dilation: Sequence[int],
+    rhs_dilation: Sequence[int],
+) -> int:
+  """Gets the number of valid spatial positions for conv_general_dilated.
+
+  Args:
+    dnums: The dimension numbers for the convolution.
+    lhs: The shape of the left-hand side of the convolution.
+    rhs: The shape of the right-hand side of the convolution.
+    out: The shape of the output of the convolution.
+    window_strides: The stride of the window along each spatial dimension.
+    padding: The padding applied to the input along each spatial dimension.
+    lhs_dilation: The dilation factor for the left-hand side along each spatial
+      dimension.
+    rhs_dilation: The dilation factor for the right-hand side along each spatial
+      dimension.
+  """
+  input_spatial_dims, kernel_spatial_dims, out_spatial_dims = (
+      dnums.lhs_spec[2:],
+      dnums.rhs_spec[2:],
+      dnums.out_spec[2:],
+  )
+
+  valid_position_counts = 1
+  # Loop over each spatial dimension and determine how many valid positions
+  # there are for each dimension.
+  for d in range(len(input_spatial_dims)):
+    valid_position_counts *= _get_spatial_valid_position_count_for_one_dim(
+        window_dim_stride=window_strides[d],
+        base_dilation=lhs_dilation[d],
+        window_dilation=rhs_dilation[d],
+        kernel_limit=rhs.shape[kernel_spatial_dims[d]],
+        input_limit=lhs.shape[input_spatial_dims[d]],
+        output_limit=out.shape[out_spatial_dims[d]],
+        padding=padding[d],
+    )
+
+  return valid_position_counts
+
+
+def _calculate_conv_flops(
+    lhs: roofline.RooflineShape,
+    rhs: roofline.RooflineShape,
+    out: roofline.RooflineShape,
+    window_strides: Sequence[int],
+    padding: Sequence[tuple[int, int]],
+    lhs_dilation: Sequence[int],
+    rhs_dilation: Sequence[int],
+    dimension_numbers: convolution.ConvGeneralDilatedDimensionNumbers,
+    batch_group_count: int,
+) -> int:
+  """Calculates roofline unfused flops for Jax's conv_general_dilated primitive.
+
+  See `jax.lax.conv_general_dilated` for details on the arguments.
+  """
+  dnums = convolution.conv_dimension_numbers(
+      lhs.shape, rhs.shape, dimension_numbers
+  )
+
+  spatial_valid_position_counts = _get_spatial_valid_position_count(
+      dnums, lhs, rhs, out, window_strides, padding, lhs_dilation, rhs_dilation
+  )
+
+  batch = lhs.shape[dnums.lhs_spec[0]]
+  num_output_features = out.shape[dnums.out_spec[1]]
+  num_input_features = rhs.shape[dnums.rhs_spec[1]]
+  num_output_batch = batch / batch_group_count
+
+  non_spatial_dims_factor = (
+      num_input_features * num_output_features * num_output_batch
+  )
+
+  fma_count = non_spatial_dims_factor * spatial_valid_position_counts
+  flops = fma_count * _FMA_FLOPS_FACTOR
+  return int(flops)
+
+
 @roofline.register_roofline(convolution.conv_general_dilated_p)
 def _conv_general_dilated_roofline(
-  ctx: roofline.RooflineRuleContext,
-  *args,
-  **kw,
+    ctx: roofline.RooflineRuleContext,
+    *args,
+    window_strides: Sequence[int],
+    padding: Sequence[tuple[int, int]],
+    lhs_dilation: Sequence[int],
+    rhs_dilation: Sequence[int],
+    dimension_numbers: convolution.ConvGeneralDilatedDimensionNumbers,
+    batch_group_count: int,
+    **kw,
 ) -> roofline.RooflineResult:
+  """Roofline for Jax's conv_general_dilated primitive.
+
+  See `jax.lax.conv_general_dilated` for details on the arguments.
+  """
   lhs, rhs = (roofline.RooflineShape.from_aval(aval) for aval in ctx.avals_in)
   out = roofline.RooflineShape.from_aval(ctx.avals_out[0])
-  # TODO(b/394648206): support computing unfused_flops for conv.
+
   return roofline.RooflineResult(
+      unfused_flops=_calculate_conv_flops(
+          lhs,
+          rhs,
+          out,
+          window_strides,
+          padding,
+          lhs_dilation,
+          rhs_dilation,
+          dimension_numbers,
+          batch_group_count,
+      ),
       unfused_hbm_bytes=(
           lhs.dtype.itemsize * lhs.size
           + rhs.dtype.itemsize * rhs.size


Test patch:
diff --git a/tests/roofline_test.py b/tests/roofline_test.py
index 98f6176c2..140beb3c6 100644
--- a/tests/roofline_test.py
+++ b/tests/roofline_test.py
@@ -28,6 +28,8 @@ from jax.sharding import PartitionSpec as P
 jax.config.parse_flags_with_absl()
 jtu.request_cpu_devices(8)
 
+_VERY_LARGE_NUMBER = 512 * 1024
+
 
 def create_inputs(
   *shardings: P,
@@ -628,7 +630,6 @@ class RooflineTest(jtu.JaxTestCase):
     expected_unfused_hbm_bytes = self._bytes_per_word * (
         expected_input_size + expected_kernel_size + expected_output_size
     )
-    # TODO(b/394648206): add subtest for unfused_flops once they are supported.
     self.assertEqual(result.unfused_hbm_bytes, expected_unfused_hbm_bytes)
 
   @jtu.parameterized.named_parameters(
@@ -641,10 +642,10 @@ class RooflineTest(jtu.JaxTestCase):
           padding="SAME_LOWER",
       ),
   )
-  def test_conv_general_dilated_padding_string_unfused_hbm_bytes(
+  def test_conv_general_dilated_padding_string(
       self, padding: str
   ):
-    input_data = jnp.zeros((1, 1, 10, 20), dtype=int)
+    input_data = jnp.zeros((1, 1, 3, 3), dtype=int)
     kernel_data = jnp.ones((1, 1, 3, 3), dtype=int)
     conv = lambda a, b: lax.conv_general_dilated(
         lhs=a, rhs=b, window_strides=(1, 1), padding=padding
@@ -652,10 +653,11 @@ class RooflineTest(jtu.JaxTestCase):
 
     _, result = roofline.roofline(conv)(input_data, kernel_data)
 
-    expected_input_size = 1 * 1 * 10 * 20
+    # Test hbm bytes.
+    expected_input_size = 1 * 1 * 3 * 3
     expected_kernel_size = 1 * 1 * 3 * 3
     # Because of same{_lower} padding, output shape should equal to input shape.
-    # This may not be true for other `{feature, batch}`_group_count`s.c
+    # This may not be true for other `{feature, batch}`_group_count`s.
     expected_output_size = expected_input_size
     # Bytes accessed is sum of inputs and output.
     expected_unfused_hbm_bytes = self._bytes_per_word * (
@@ -663,7 +665,21 @@ class RooflineTest(jtu.JaxTestCase):
     )
     self.assertEqual(result.unfused_hbm_bytes, expected_unfused_hbm_bytes)
 
-  def test_conv_general_dilated_padding_string_valid_unfused_hbm_bytes(self):
+    # Test flops.
+    # For spatial_valid_position_counts, we have 3x3 output with the following
+    # flops for each element:
+    #    4 6 4
+    #    6 9 6
+    #    4 6 4
+    # Non_spatial_dims_factor = 1 because `{batch, feature}_group_count` are
+    # both equal to 1.
+    # Each FMA is 2 flops.
+    self.assertEqual(
+        result.unfused_flops,
+        2 * (4 + 6 + 4 + 6 + 9 + 6 + 4 + 6 + 4),
+    )
+
+  def test_conv_general_dilated_padding_string_valid(self):
     input_data = jnp.zeros((1, 1, 10, 20), dtype=int)
     kernel_data = jnp.ones((1, 1, 3, 3), dtype=int)
     conv = lambda a, b: lax.conv_general_dilated(
@@ -681,12 +697,90 @@ class RooflineTest(jtu.JaxTestCase):
         * self.get_conv_output_dim(10, 3, 0, 0, 1)
         * self.get_conv_output_dim(20, 3, 0, 0, 1)
     )
+
     # Bytes accessed is sum of inputs and output.
     expected_unfused_hbm_bytes = self._bytes_per_word * (
         expected_input_size + expected_kernel_size + expected_output_size
     )
     self.assertEqual(result.unfused_hbm_bytes, expected_unfused_hbm_bytes)
 
+    # Output shape is [1x1x8x18] and each output element requires (3x3) FMAs,
+    # and each FMA is 2 flops.
+    self.assertEqual(
+        result.unfused_flops, 2 * expected_output_size * 3 * 3
+    )
+
+
+  @jtu.parameterized.named_parameters(
+      dict(
+          testcase_name="padding",
+          input_spatial_dim=1,
+          window_strides=[1],
+          padding=[(_VERY_LARGE_NUMBER - 1, _VERY_LARGE_NUMBER - 1)],
+          lhs_dilation=[1],
+      ),
+      dict(
+          testcase_name="input",
+          input_spatial_dim=_VERY_LARGE_NUMBER,
+          window_strides=[_VERY_LARGE_NUMBER - 1],
+          padding=[(0, 0)],
+          lhs_dilation=[_VERY_LARGE_NUMBER],
+      ),
+  )
+  def test_conv_general_dilated_flops_very_large(
+      self, input_spatial_dim, window_strides, padding, lhs_dilation
+  ):
+    input_data = jnp.zeros((1, 1, input_spatial_dim), dtype=int)
+    kernel_data = jnp.ones((1, 1, _VERY_LARGE_NUMBER), dtype=int)
+    conv = lambda a, b: lax.conv_general_dilated(
+        lhs=a,
+        rhs=b,
+        window_strides=window_strides,
+        padding=padding,
+        lhs_dilation=lhs_dilation,
+    )
+    _, result = roofline.roofline(conv)(input_data, kernel_data)
+
+    self.assertEqual(result.unfused_flops, 2 * _VERY_LARGE_NUMBER)
+
+  def test_conv_general_dilated_flops_feature_group_count(self):
+    feature_group_count = 120
+    input_data = jnp.zeros((1, feature_group_count, 10, 20), dtype=int)
+    kernel_data = jnp.ones((feature_group_count, 1, 3, 3), dtype=int)
+    conv = lambda a, b: lax.conv_general_dilated(
+        lhs=a,
+        rhs=b,
+        window_strides=(1, 1),
+        padding=((0, 0), (0, 0)),
+        feature_group_count=feature_group_count,
+    )
+    _, result = roofline.roofline(conv)(input_data, kernel_data)
+
+    # Output shape is [1x120x8x18] and each output element requires (3x3)
+    # FMAs and one FMA is 2 flops.
+    self.assertEqual(
+        result.unfused_flops, 2 * 120 * 8 * 18 * 3 * 3
+    )
+
+  def test_conv_general_dilated_flops_batch_group_count(self):
+    batch_group_count = 120
+    input_data = jnp.zeros((batch_group_count, 1, 10, 20), dtype=int)
+    kernel_data = jnp.ones((batch_group_count, 1, 3, 3), dtype=int)
+    conv = lambda a, b: lax.conv_general_dilated(
+        lhs=a,
+        rhs=b,
+        window_strides=(1, 1),
+        padding=((0, 0), (0, 0)),
+        batch_group_count=batch_group_count,
+    )
+    _, result = roofline.roofline(conv)(input_data, kernel_data)
+
+    # Output shape is [120x1x8x18] and each output element requires (3x3)
+    # FMAs and one FMA is 2 flops.
+    self.assertEqual(
+        result.unfused_flops, 2 * 120 * 8 * 18 * 3 * 3
+    )
+
   def test_reduce_sum_no_axis(self):
     _, result = roofline.roofline(lambda x: jnp.sum(x))(jnp.zeros((11, 4)))
     self.assertEqual(result.unfused_flops, 11 * 4 - 1)