metal lowbit kernels: replace ATEN by USE_ATEN

Differential Revision: D65253390 Pull Request resolved: pytorch#1263
sunjiweiswift · Nov 25, 2024 · e6dc429 · e6dc429
1 parent b6bcd08
commit e6dc429
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/torchao/experimental/kernels/mps/codegen/gen_metal_shader_lib.py b/torchao/experimental/kernels/mps/codegen/gen_metal_shader_lib.py
@@ -28,7 +28,7 @@
  * This file is generated by gen_metal_shader_lib.py
  */
 
-#ifdef ATEN
+#ifdef USE_ATEN
 using namespace at::native::mps;
 #else
 #include <torchao/experimental/kernels/mps/src/OperationUtils.h>

diff --git a/torchao/experimental/kernels/mps/src/lowbit.h b/torchao/experimental/kernels/mps/src/lowbit.h
@@ -17,7 +17,7 @@
 #include <fstream>
 #include <sstream>
 
-#ifdef ATEN
+#ifdef USE_ATEN
 #include <ATen/native/mps/OperationUtils.h>
 using namespace at::native::mps;
 inline void finalize_block(MPSStream* mpsStream) {}

diff --git a/torchao/experimental/ops/mps/setup.py b/torchao/experimental/ops/mps/setup.py
@@ -16,7 +16,7 @@
             name="torchao_mps_ops",
             sources=["register.mm"],
             include_dirs=[os.getenv("TORCHAO_ROOT")],
-            extra_compile_args=["-DATEN=1"],
+            extra_compile_args=["-DUSE_ATEN=1"],
         ),
     ],
     cmdclass={"build_ext": BuildExtension},