p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] test/test.py (python)
# import jax
# with jax.profiler.trace("tensorboard"):
  # key = jax.random.key(0)
  # x = jax.random.normal(key, (1024, 1024))
  # y = x @ x
  # y.block_until_ready()

import datetime
import jax.numpy as jnp
import jax

MATRIX_DIM = 32768
STEPS = 10

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

num_bytes = A.size * 4
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

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}")
[tab] test/test_jax.py (python)
# import jax
# with jax.profiler.trace("tensorboard"):
  # key = jax.random.key(0)
  # x = jax.random.normal(key, (1024, 1024))
  # y = x @ x
  # y.block_until_ready()

import datetime
import jax.numpy as jnp
import jax

MATRIX_DIM = 32768*2
STEPS = 10

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

num_bytes = A.size * 4
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

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}")
[content] test/test_torch.py (python)
import torch

MATRIX_DIM = 32768 * 2
STEPS = 10

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

A = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)
B = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)

num_bytes = A.numel() * A.element_size()
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(X, Y):
  return X @ Y

# warmup
with torch.no_grad():
  _ = matmul(A, B)
  if device.type == 'cuda':
    torch.cuda.synchronize()

start_time = datetime.datetime.now()
with torch.no_grad():
  for i in range(STEPS):
    C = A @ B
    if device.type == 'cuda':
      torch.cuda.synchronize()
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}")
[content] test/test_torch.py (python)
import datetime
import torch

MATRIX_DIM = 32768 * 2
STEPS = 10

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

A = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)
B = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)

num_bytes = A.numel() * A.element_size()
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(X, Y):
  return X @ Y

# warmup
with torch.no_grad():
  _ = matmul(A, B)
  if device.type == 'cuda':
    torch.cuda.synchronize()

start_time = datetime.datetime.now()
with torch.no_grad():
  for i in range(STEPS):
    C = A @ B
    if device.type == 'cuda':
      torch.cuda.synchronize()
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}")
[content] test/test_torch.py (python)
# pyright: reportMissingImports=false
import datetime
import torch

MATRIX_DIM = 32768 * 2
STEPS = 10

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

A = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)
B = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)

num_bytes = A.numel() * A.element_size()
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(X, Y):
  return X @ Y

# warmup
with torch.no_grad():
  _ = matmul(A, B)
  if device.type == 'cuda':
    torch.cuda.synchronize()

start_time = datetime.datetime.now()
with torch.no_grad():
  for i in range(STEPS):
    C = A @ B
    if device.type == 'cuda':
      torch.cuda.synchronize()
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}")