Wrong metrics for FLOP utilization and TensorCore usage when using bfloat16

[emergenz]

The following program is a matrix multiplication of two matrices in BF16 on an H100. We achieve a rough throughput of 600 TFLOP/s (which is impossible without using TensorCores). BF16 matrix multiplication is TensorCore-eligible on the H100. The GPU kernel stats page of the JAX profiler incorrectly states that the op is not TensorCore-eligible and that no TensorCores are used. The framework op stats page also incorrectly identifies the op as not being TensorCore-eligible.

However, again, such high throughput is impossible without using TensorCores (I stated the throughput manually measured by myself, which corresponds to the one listed in the graph viewer).

On a slightly different note, for the FLOPS utilization it is unclear to me whether the peak FLOP/s used for the calculation are correct (which should be different for BF16 and TF32).

[emergenz]

I just realized that I didn't append the actual program to this issue. My bad. It was exactly the program in #1704 but with the inputs cast to BF16.

import datetime
import jax.numpy as jnp
import jax

MATRIX_DIM = 32768
STEPS = 10

A = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.bfloat16)
B = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.bfloat16)

num_bytes = A.size * 2
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(A, B):
  return A @ B

matmul(A, B) # warmup

with jax.profiler.trace("tensorboard"):
  start_time = datetime.datetime.now()
  for i in range(STEPS):
    C = A @ B
    C.block_until_ready()
  end_time = datetime.datetime.now()

average_time_per_step = (end_time - start_time).total_seconds() / STEPS

print(f"{average_time_per_step}, teraflops per second: {total_num_flops / average_time_per_step / 1e12}, gigabytes per second: {total_num_bytes_crossing_to_hbm / average_time_per_step / 1e9}")