add a benchmark for casting a tensor to MX across dim0 and dim1

[vkuzo]

Summary:

This is useful for training, extracting into a benchmark so we can
optimize.

Test Plan:

TORCH_LOGS_FORMAT=short TORCH_LOGS=aot_graphs,output_code python benchmarks/float8/profile_lowp_training.py ~/local/tmp/20250223_test --mx_recipe_name mxfp8_emulated --experiment_filter lowp --mode_filter cast_only_dim0_dim1
// output: https://gist.github.com/vkuzo/a4e13bac7fc8ca3af10bfd5483b85b33
// currently we see two kernels, one per dim

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> add a benchmark for casting a tensor to MX across dim0 and dim1 #1787

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1787

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 5 Pending

As of commit 2a4b6ff with merge base 8d110bf ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test (macos-14) (gh)
  [ FAILED ] test_linear_8bit_act_xbit_weight.Kleidi_dotprod_1x8x32_m101xn34xk128xg64

This comment was automatically generated by Dr. CI and updates every 15 minutes.