cuDNN depthwise conv example

conv.py

import torch
import time

torch.backends.cudnn.benchmark = True

iters = 10

conv = torch.nn.Conv2d(64, 64, 3, 3, groups=64, dtype=torch.half, device='cuda')
convb = torch.nn.Conv2d(64, 64, 3, 3, groups=64, dtype=torch.bfloat16, device='cuda')
data = torch.randn(16, 64, 1024, 1024, dtype=torch.half, device='cuda')
datab = torch.randn(16, 64, 1024, 1024, dtype=torch.bfloat16, device='cuda')

# half
# warmup
out = conv(data)
torch.cuda.synchronize()
t1 = time.time()
for _ in range(iters):
  out = conv(data)
torch.cuda.synchronize()
t2 = time.time()
print(f"half took {(t2-t1)/iters} per iteration")

# bfloat16
# warmup
outb = convb(datab)
torch.cuda.synchronize()
t1 = time.time()
for _ in range(iters):
  outb = convb(datab)
torch.cuda.synchronize()
t2 = time.time()
print(f"bfloat16 took {(t2-t1)/iters} per iteration")

Changing 4 lines (the torch.nn.functional.conv2d's) produces slower times for both precisions and produces the bfloat16 <> float16 mismatch:

import torch
import time

torch.backends.cudnn.benchmark = True

iters = 10

conv = torch.nn.Conv2d(64, 64, 3, 3, groups=64, dtype=torch.half, device='cuda')
convb = torch.nn.Conv2d(64, 64, 3, 3, groups=64, dtype=torch.bfloat16, device='cuda')
data = torch.randn(16, 64, 1024, 1024, dtype=torch.half, device='cuda')
datab = torch.randn(16, 64, 1024, 1024, dtype=torch.bfloat16, device='cuda')

# half
# warmup
out = torch.nn.functional.conv2d(data, conv.weight, groups=64)
torch.cuda.synchronize()
t1 = time.time()
for _ in range(iters):
  out = torch.nn.functional.conv2d(data, conv.weight, groups=64)
torch.cuda.synchronize()
t2 = time.time()
print(f"half took {(t2-t1)/iters} per iteration")

# bfloat16
# warmup
outb = torch.nn.functional.conv2d(datab, convb.weight, groups=64)
torch.cuda.synchronize()
t1 = time.time()
for _ in range(iters):
  outb = torch.nn.functional.conv2d(datab, convb.weight, groups=64)
torch.cuda.synchronize()
t2 = time.time()
print(f"bfloat16 took {(t2-t1)/iters} per iteration")