From d509832426ddaaa6263640b5040962e334fdbbe3 Mon Sep 17 00:00:00 2001
From: Tri Dao <tridpq@gmail.com>
Date: Thu, 12 Jan 2023 22:15:41 -0800
Subject: [PATCH] [Compilation] Add _NO_HALF2 flags to be consistent with
 Pytorch

https://github.com/pytorch/pytorch/blob/eb7b89771e989f832147b8f28f62c0b43d5d3c12/cmake/Dependencies.cmake#L1693
---
 setup.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/setup.py b/setup.py
index 89b5d27..3fb7666 100644
--- 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",