[reland][ROCm] preshuffled weight mm

[jeffdaily]

Adds SwizzleTensor subclass that wraps a Tensor and reorders the contents to be suitable for HIPBLASLT_ORDER_COL16_4R8. SwizzleTensor intercepts torch.mm and replaces with custom calls to hipblaslt.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 27 New Failures

As of commit 53000ac with merge base a81322e ():

NEW FAILURES - The following jobs have failed:

• Build Linux Wheels / build / build-manywheel-py3_9-cpu (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-cuda11_8 (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-cuda12_6 (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-cuda12_8 (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-rocm6_2_4 (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-rocm6_3 (gh)
  RuntimeError: Error compiling objects for extension
• Build Linux Wheels / build / build-manywheel-py3_9-xpu (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cpu (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda11_8 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_ao__3.9_cu118_x86_64
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda12_6 (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda12_8 (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-rocm6_2_4 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_ao__3.9_rocm6.2.4_x86_64
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-rocm6_3 (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-xpu (gh)
• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run Float8 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio --index-url https://download.pytor... / linux-job (gh)
  RuntimeError: Command docker exec -t 5a6bc401e5ecdd33280d270d0e2d8ab62bdb497c9b71461a487ddb84ea7bcedc /exec failed with exit code 1
• Run Float8 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 7a98692182886f72f9eec1e07e5934640ffa43fbca24b7eaefa98304c7739f26 /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 86c32fde8a1b6b3a3ece8aea9c61af6d1aab984b4c5b4beff7132a7a989ba9e4 /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t aec50e8cd5b45e6864123e19b973a42914b1249d6b9025e4c887032ede6c8506 /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 192f6e074d652ff59900204401d6c4928979b73c4ac228edf80efcecb70b739d /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 8bf4c9d40a648019e01a51f9dbc254ab59ba41e6ad8d4dcff882f3f0865d6b6a /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 34ea30e707ba81906fcd8fe9553887275596c16116cbbdbeec859e21b67ba0fc /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  RuntimeError: Command docker exec -t a55bd3cd4741d731f4f8c80cb3962746198a479b8bd37f9c59249e87061ec89e /exec failed with exit code 1
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t c690c12a7bbb556f73ccab5167f08fde8c727101a4180a0958a294aeee5e6a7e /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 8efcbaa8dbf9f36ded9c0b93af7b1d53f08c37e393bf95d69d2e2d471e898aad /exec failed with exit code 1
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  ninja: build stopped: subcommand failed
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  ninja: build stopped: subcommand failed

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

[jeffdaily]

@mxz297 @jerryzh168 please re-review, kick of CI, thanks.

[mxz297]

@jeffdaily "test-mps-ops" still seems to be failing to compile with

  /Users/ec2-user/runner/_work/ao/ao/torchao/csrc/rocm/swizzle/swizzle.cpp:1:10: fatal error: 'hip/hip_runtime.h' file not found
  #include <hip/hip_runtime.h>
           ^~~~~~~~~~~~~~~~~~~

I wonder if we should just guard the whole source file under #if USE_ROCM