From 12d6418e172b5a9d64157c4671c74bf61fac3ce2 Mon Sep 17 00:00:00 2001
From: Qubitium-ModelCloud <qubitium@modelcloud.ai>
Date: Sat, 12 Oct 2024 01:01:41 -0700
Subject: [PATCH] [BUILD] update compile flags (#428)

* [WIP] [BUILD] Test new compile flags

* Update setup.py

* Update setup.py
---
 setup.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index a945c4a9..5e116f73 100644
--- a/setup.py
+++ b/setup.py
@@ -5,8 +5,6 @@
 
 from setuptools import find_packages, setup
 
-os.environ["CC"] = "g++"
-os.environ["CXX"] = "g++"
 os.environ["BUILD_CUDA_EXT"] = "1"
 
 TORCH_CUDA_ARCH_LIST = os.environ.get("TORCH_CUDA_ARCH_LIST")
@@ -112,17 +110,29 @@
     extra_compile_args = {
         "cxx": [
             "-O3",
+            "-std=c++17", 
+            "-fopenmp", 
+            "-lgomp", 
+            "-DENABLE_BF16"
             "-Wno-switch-bool",
         ],
         "nvcc": [
             "-O3",
             "-std=c++17",
+            "-DENABLE_BF16",
+            "-U__CUDA_NO_HALF_OPERATORS__",
+            "-U__CUDA_NO_HALF_CONVERSIONS__",
+            "-U__CUDA_NO_HALF2_OPERATORS__",
+            "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+            "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+            "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
             "--threads",
             "4",
             "-Xfatbin",
             "-compress-all",
             "-diag-suppress=179,39,186",
-             "--use_fast_math",
+            "--use_fast_math",
         ],
     }