[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

mlx/mlx/backend/metal/kernels/conv.metal

Lines 587 to 588 in 2e7c02d

	instantiate_winograd_conv_2d(float32, float);
	instantiate_winograd_conv_2d(float16, half);

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:

mlx/mlx/backend/metal/conv.cpp

Lines 649 to 655 in 2e7c02d

	// Direct to winograd conv
	if (!flip && is_stride_one && is_kdil_one && is_idil_one &&
	conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
	conv_params.C % 32 == 0 && conv_params.O % 32 == 0 &&
	(channels_large || (channels_med && inp_large))) {
	return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
	}

(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