Keep the winograd gemm at full float32 precision

[erwinzhang7]

A float32 conv2d on hardware with neural accelerators returns roughly float16 accuracy, but only
when the channel counts are multiples of 32.

conv.cpp sends 3x3 stride 1 convs to winograd when C % 32 == 0 and O % 32 == 0 (plus size
conditions). Winograd runs its batched gemm through steel_matmul, which takes the nax path for
float32 whenever MLX_ENABLE_TF32 is set, and that is the default. The winograd transforms then
amplify what tf32 rounded away. The implicit gemm path never calls steel_matmul, so it is
unaffected, which is why the same conv is accurate at 336 channels and not at 352.

N=4, 64x64, 3x3, padding 1, C = O, error against a cpu float32 reference:

dtype	implicit gemm	winograd	ratio
float32	0.00085	4.15259	4859.7x
float16	0.14026	5.13802	36.6x
bfloat16	0.98453	42.59959	43.3x

Setting MLX_ENABLE_TF32=0 drops the float32 winograd error to 0.00693, a 600x change, and the
same shape on an M4 Pro, which has no neural accelerators, measures 0.00693 as well. Both point
at the same place, so this makes winograd ask for full float32 rather than depending on a global
that is about matmul.

steel_matmul_axpby decides use_nax from env::enable_tf32() internally, so this adds an
allow_tf32 parameter defaulted to true and threads it to that decision. Every existing caller
is unchanged; winograd passes false. The flag only has an effect for float32 inputs, since the
condition it guards is already false for every other dtype.

The float16 and bfloat16 rows above are a separate problem and are not addressed here. Those come
from the winograd intermediates being allocated in the input dtype, and fixing them means
doubling the working set, which is the opposite of what #4102 is trying to do.

Cost

N=8, 64x64, C = O = 352, calibrated, min of 7.

	before	after
float32 winograd	2.124 ms	2.654 ms
float16 winograd	1.349 ms	1.333 ms
bfloat16 winograd	1.324 ms	1.328 ms
implicit gemm, C=336	5.312 ms	5.313 ms

25% on the float32 winograd path and nothing anywhere else. It stays about 2x faster than the
implicit gemm the same conv falls to with an unaligned channel count.

Test

mlx_tests sets MLX_ENABLE_TF32=0 for the whole suite and enable_tf32() reads the
environment once into a static, so a test in the suite cannot see this. The new test runs the
conv in a subprocess with tf32 on and compares an aligned conv against an unaligned one. It
reports 2.38 before this change and 0.0023 after.

test_conv, test_blas, test_ops, test_autograd and test_nn all pass, which is worth
saying because matmul.h has a lot of callers.

This touches the steel_matmul call that #4102 moves, so whichever lands second needs a one line
rebase.

[JasonHonKL]

curious why need to use sub process for this test ?

[erwinzhang7]

mlx_tests.py sets MLX_ENABLE_TF32=0 at import, before mlx.core is imported, for the whole
suite:

# Use regular fp32 precision for tests
os.environ["MLX_ENABLE_TF32"] = "0"

and enable_tf32() reads the environment once into a static:

inline bool enable_tf32() {
  static bool enable_tf32_ = get_var("MLX_ENABLE_TF32", 1);
  return enable_tf32_;
}

So by the time any test runs, tf32 is off and the value is fixed for the life of the process.
Setting os.environ inside the test does nothing, and this change is a no-op with tf32 off, so an
in-process test would vacuously pass.

The subprocess runs the conv with MLX_ENABLE_TF32=1, which is the default a user gets. It
reports 2.38 before the change and 0.0023 after.

The same two facts mean the suite currently has no coverage of the default
tf32 configuration at all, which is why a 4859x accuracy difference on float32 conv went
unnoticed.