[Compilation] Add _NO_HALF2 flags to be consistent with Pytorch

https://github.com/pytorch/pytorch/blob/eb7b89771e989f832147b8f28f62c0b43d5d3c12/cmake/Dependencies.cmake#L1693

[Compilation] Add _NO_HALF2 flags to be consistent with Pytorch
https://github.com/pytorch/pytorch/blob/eb7b89771e989f832147b8f28f62c0b43d5d3c12/cmake/Dependencies.cmake#L1693
d5098324 · Tri Dao · f95c2fc1 · d5098324
Commit d5098324 authored Jan 12, 2023 by Tri Dao
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

setup.py setup.py +2 -0

No files found.
--- a/setup.py
+++ b/setup.py
@@ -136,6 +136,8 @@ ext_modules.append(
                    "-std=c++17",
                    "-U__CUDA_NO_HALF_OPERATORS__",
                    "-U__CUDA_NO_HALF_CONVERSIONS__",
+                    "-U__CUDA_NO_HALF2_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                    "--expt-relaxed-constexpr",
                    "--expt-extended-lambda",
                    "--use_fast_math",