[BUG] Unable to load kernel winograd_conv_2d_weight_transform_bfloat16_bc32

[kaeru-shigure]

Describe the bug

libc++abi: terminating due to uncaught exception of type std::runtime_error: [metal::Device] Unable to load kernel winograd_conv_2d_weight_transform_bfloat16_bc32

https://github.com/ml-explore/mlx/blob/2e7c02d5cdb173c777e42128c1590e7d86dc9a55/mlx/backend/metal/kernels/conv.metal#L587-L588

Maybe there is a lack of implementation around here, but I'm not familiar with METAL, so I couldn't fix it.

To Reproduce

Include code snippet

# I couldn't think of a minimal reproduction code.
# But I think the cause is simply unimplemented.

Expected behavior works perfect

Desktop (please complete the following information):

• OS Version: macOS 14.4.1
• Version: 0.10.0

[awni]

CC @jagrit06 any reason that one is not instantiated for bfloat?

Also @kaeru-shigure how did you come across this? Could you share the program you ran?

[kaeru-shigure]

The winograd algorithm is selected when Conv2d satisfies the following conditions and when this is a bfloat16 weight, an unimplemented error occurs:

https://github.com/ml-explore/mlx/blob/2e7c02d5cdb173c777e42128c1590e7d86dc9a55/mlx/backend/metal/conv.cpp#L649-L655

(commenting out exactly this section seems to be a workaround). I cannot provide a reproduction code because the conditions are too complex. sorry.

[jagrit06]

I’ll look into the Winifred code for bf16