Metal Shader Compilation Error: bfloat type not defined (commit 312e264)

[fuzzycuffs]

Metal Shader Compilation Error: bfloat type not defined

Environment

• macOS: 26.2 (Tahoe)
• Hardware: Apple M4 Pro
• Metal: Metal 4
• PyTorch: 2.10.0
• mps-deform-conv: 0.2.1 (commit 312e264)
• Python: 3.12

Issue

Commit 312e264 ("Add native BF16 kernel...") introduced a deformable_im2col_bf16 kernel that uses the type bfloat, but this type is not defined anywhere in the Metal shader.

The shader compilation fails with:

Failed to compile Metal library: program_source:249:18: error: unknown type name 'bfloat'; did you mean 'float'?
    device const bfloat* input       [[buffer(0)]],
                 ^~~~~~
                 float

Root Cause

Looking at csrc/deform_conv2d_mps.mm lines 249-336, the BF16 kernel uses bfloat directly:

kernel void deformable_im2col_bf16(
    device const bfloat* input       [[buffer(0)]],
    device const bfloat* offset      [[buffer(1)]],
    ...

But Metal does not have a native bfloat type. Metal supports:

• float (32-bit)
• half (16-bit IEEE FP16)

BFloat16 is not part of the Metal Shading Language. It would need to be emulated or defined as a typedef.

Reproduction

from mps_deform_conv import deform_conv2d
import torch

device = torch.device('mps')
x = torch.randn(1, 64, 32, 32, device=device, dtype=torch.float32)
offset = torch.randn(1, 2*9, 32, 32, device=device, dtype=torch.float32)
weight = torch.randn(64, 64, 3, 3, device=device, dtype=torch.float32)
mask = torch.randn(1, 9, 32, 32, device=device, dtype=torch.float32)

# Fails even with float32 because dispatch logic selects BF16 path
out = deform_conv2d(x, offset, weight, None, (1,1), (1,1), (1,1), mask)

Suggested Fix

Option 1: Remove BF16 kernel entirely (simplest):

// Don't dispatch to bf16 path

Option 2: Use half instead since MPS doesn't support true BF16:

kernel void deformable_im2col_bf16(
    device const half* input  // Use half instead of bfloat
    ...

Option 3: Add typedef (if Metal 3.1+ supports a bfloat type):

// At top of shader
typedef half bfloat;  // Or appropriate type

Notes

The error occurs even with float32 input tensors, suggesting the dispatch logic in getComputePipeline() is incorrectly selecting the BF16 path when it shouldn't.

Thanks for the great work on this library! Looking forward to using it with BasicVSR++ for video restoration.

[imperatormk]

Thank you for this report, on my Metal 4 and 26.3 all this works well but I added a fallback so please let me know if has been resolved now!