From be5d4d6107b4e658c745b008923b58f413c0de24 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 22 Oct 2025 08:42:52 +0000
Subject: [PATCH] python3Packages.torchao: fix build on darwin

---
 .../python-modules/torchao/default.nix        | 61 +++++++++++++++++++
 .../torchao/use-llvm-openmp.patch             | 13 ++++
 .../torchao/use-system-cpuinfo.patch          | 22 +++++++
 3 files changed, 96 insertions(+)
 create mode 100644 pkgs/development/python-modules/torchao/use-llvm-openmp.patch
 create mode 100644 pkgs/development/python-modules/torchao/use-system-cpuinfo.patch

diff --git a/pkgs/development/python-modules/torchao/default.nix b/pkgs/development/python-modules/torchao/default.nix
index e92454771c3d..2af60a2115b0 100644
--- a/pkgs/development/python-modules/torchao/default.nix
+++ b/pkgs/development/python-modules/torchao/default.nix
@@ -3,10 +3,18 @@
   stdenv,
   buildPythonPackage,
   fetchFromGitHub,
+  replaceVars,
 
   # build-system
   setuptools,
 
+  # nativeBuildInputs
+  cmake,
+
+  # buildInputs
+  cpuinfo,
+  llvmPackages,
+
   # dependencies
   torch,
 
@@ -34,10 +42,33 @@ buildPythonPackage rec {
     hash = "sha256-L9Eoul7Nar/+gS44+hA8JbfxCgkMH5xAMCleggAZn7c=";
   };
 
+  patches = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    ./use-system-cpuinfo.patch
+    (replaceVars ./use-llvm-openmp.patch {
+      inherit (llvmPackages) openmp;
+    })
+  ];
+
   build-system = [
     setuptools
   ];
 
+  nativeBuildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    cmake
+  ];
+  dontUseCmakeConfigure = true;
+
+  buildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    cpuinfo
+  ];
+
+  propagatedBuildInputs = lib.optionals stdenv.hostPlatform.isDarwin [
+    # Otherwise, torch will fail to include `omp.h`:
+    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error
+    # OpenMP support not found.
+    llvmPackages.openmp
+  ];
+
   dependencies = [
     torch
   ];
@@ -90,6 +121,36 @@ buildPythonPackage rec {
     "test_save_load_int4woqtensors_2_cpu"
     "test_save_load_int8woqtensors_0_cpu"
     "test_save_load_int8woqtensors_1_cpu"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
+    # AssertionError: Scalars are not equal!
+    "test_comm"
+    "test_fsdp2"
+    "test_fsdp2_correctness"
+    "test_precompute_bitnet_scale"
+    "test_qlora_fsdp2"
+    "test_uneven_shard"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    # RuntimeError: No packed_weights_format was selected
+    "TestIntxOpaqueTensor"
+    "test_accuracy_kleidiai"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
+    # Flaky: [gw0] node down: keyboard-interrupt
+    "test_int8_weight_only_quant_with_freeze_0_cpu"
+    "test_int8_weight_only_quant_with_freeze_1_cpu"
+    "test_int8_weight_only_quant_with_freeze_2_cpu"
+  ];
+
+  disabledTestPaths = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    # Require unpackaged 'coremltools'
+    "test/prototype/test_groupwise_lowbit_weight_lut_quantizer.py"
+
+    # AttributeError: '_OpNamespace' 'mkldnn' object has no attribute '_is_mkldnn_acl_supported'
+    "test/quantization/pt2e/test_arm_inductor_quantizer.py"
+    "test/quantization/pt2e/test_x86inductor_fusion.py"
+    "test/quantization/pt2e/test_x86inductor_quantizer.py"
   ];
 
   meta = {
diff --git a/pkgs/development/python-modules/torchao/use-llvm-openmp.patch b/pkgs/development/python-modules/torchao/use-llvm-openmp.patch
new file mode 100644
index 000000000000..b399b347da39
--- /dev/null
+++ b/pkgs/development/python-modules/torchao/use-llvm-openmp.patch
@@ -0,0 +1,13 @@
+diff --git a/torchao/csrc/cpu/shared_kernels/Utils.cmake b/torchao/csrc/cpu/shared_kernels/Utils.cmake
+index be7004784..0e1a2ed0e 100644
+--- a/torchao/csrc/cpu/shared_kernels/Utils.cmake
++++ b/torchao/csrc/cpu/shared_kernels/Utils.cmake
+@@ -21,7 +21,7 @@ function(target_link_torchao_parallel_backend target_name torchao_parallel_backe
+         target_link_libraries(${target_name} PRIVATE "${TORCH_LIBRARIES}")
+ 
+         target_compile_definitions(${target_name} PRIVATE TORCHAO_PARALLEL_ATEN=1 AT_PARALLEL_OPENMP=1 INTRA_OP_PARALLEL=1)
+-        target_link_libraries(${target_name} PRIVATE ${TORCH_INSTALL_PREFIX}/lib/libomp${CMAKE_SHARED_LIBRARY_SUFFIX})
++        target_link_libraries(${target_name} PRIVATE @openmp@/lib/libomp${CMAKE_SHARED_LIBRARY_SUFFIX})
+ 
+     elseif(TORCHAO_PARALLEL_BACKEND_TOUPPER STREQUAL "EXECUTORCH")
+         message(STATUS "Building with TORCHAO_PARALLEL_BACKEND=TORCHAO_PARALLEL_EXECUTORCH")
diff --git a/pkgs/development/python-modules/torchao/use-system-cpuinfo.patch b/pkgs/development/python-modules/torchao/use-system-cpuinfo.patch
new file mode 100644
index 000000000000..25b08a6dcde9
--- /dev/null
+++ b/pkgs/development/python-modules/torchao/use-system-cpuinfo.patch
@@ -0,0 +1,22 @@
+diff --git a/torchao/csrc/cpu/CMakeLists.txt b/torchao/csrc/cpu/CMakeLists.txt
+index aaea27ec7..53dee80fd 100644
+--- a/torchao/csrc/cpu/CMakeLists.txt
++++ b/torchao/csrc/cpu/CMakeLists.txt
+@@ -49,16 +49,7 @@ if (NOT TARGET cpuinfo)
+     add_compile_options(-Wno-unused-function -Wno-unused-variable)
+ 
+     # set(CMAKE_POLICY_VERSION_MINIMUM 3.5)
+-    include(FetchContent)
+-    set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "" FORCE)
+-    set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "" FORCE)
+-    set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+-    FetchContent_Declare(cpuinfo
+-        GIT_REPOSITORY https://github.com/pytorch/cpuinfo.git
+-        GIT_TAG c61fe919607bbc534d7a5a5707bdd7041e72c5ff
+-    )
+-    FetchContent_MakeAvailable(
+-        cpuinfo)
++    find_package(cpuinfo REQUIRED)
+ 
+     cmake_policy(POP)
+ endif()