From 93ca6b78058fc7bf5d5a94b3d6704aac7b8018c1 Mon Sep 17 00:00:00 2001
From: Leonard Lausen <lausen@amazon.com>
Date: Fri, 1 May 2020 18:32:30 +0000
Subject: [PATCH] Update to thrust 1.9.8 on Windows

---
 CMakeLists.txt      |  3 +-
 ci/build_windows.py | 79 ++++++++++++++++++++++++++-------------------
 2 files changed, 47 insertions(+), 35 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index faa48e8bc2c3..998b97b192df 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -160,7 +160,8 @@ if(MSVC)
   add_definitions(-DNNVM_EXPORTS)
   add_definitions(-DNOMINMAX)
   set(CMAKE_C_FLAGS "/MP")
-  set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} /bigobj")
+  set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} ${CMAKE_CXX_FLAGS} /bigobj")
+  message("CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")
 else()
   include(CheckCXXCompilerFlag)
   set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wno-sign-compare")
diff --git a/ci/build_windows.py b/ci/build_windows.py
index 20be6ee51443..2d44f1a4e074 100755
--- a/ci/build_windows.py
+++ b/ci/build_windows.py
@@ -31,6 +31,7 @@
 import tempfile
 import time
 import zipfile
+import requests
 from distutils.dir_util import copy_tree
 from enum import Enum
 from subprocess import check_call, call
@@ -127,7 +128,6 @@ class BuildFlavour(Enum):
         '-DUSE_LAPACK=ON '
         '-DUSE_DIST_KVSTORE=OFF '
         '-DMXNET_CUDA_ARCH="5.2" '
-        '-DCMAKE_CXX_FLAGS="/FS /MD /O2 /Ob2" '
         '-DUSE_MKL_IF_AVAILABLE=OFF '
         '-DCMAKE_BUILD_TYPE=Release')
 
@@ -144,7 +144,6 @@ class BuildFlavour(Enum):
         '-DUSE_DIST_KVSTORE=OFF '
         '-DMXNET_CUDA_ARCH="5.2" '
         '-DUSE_MKLDNN=ON '
-        '-DCMAKE_CXX_FLAGS="/FS /MD /O2 /Ob2" '
         '-DCMAKE_BUILD_TYPE=Release')
 
 }
@@ -155,44 +154,56 @@ def windows_build(args):
 
     path = args.output
 
-    # cuda thrust + VS 2019 is flaky: try multiple times if fail
-    MAXIMUM_TRY = 5
-    build_try = 0
+    if os.path.exists(path):
+        shutil.rmtree(path)
+    os.makedirs(path, exist_ok=True)
 
-    while build_try < MAXIMUM_TRY:
-        if os.path.exists(path):
-            shutil.rmtree(path)
-        os.makedirs(path, exist_ok=True)
+    mxnet_root = get_mxnet_root()
+    logging.info("Found MXNet root: {}".format(mxnet_root))
 
-        mxnet_root = get_mxnet_root()
-        logging.info("Found MXNet root: {}".format(mxnet_root))
+    if 'GPU' in args.flavour:
+        # Get Thrust version to be shipped in Cuda 11, due to flakyness of
+        # older Thrust versions with MSVC 19 compiler
+        thrust_path = "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.2\\include\\thrust"
+        # if os.path.exists(thrust_path):
+        #     logging.info("Deleteting Cuda 10.2 Thrust")
+        #     shutil.rmtree(thrust_path)
 
         with remember_cwd():
-            os.chdir(path)
-            cmd = "\"{}\" && cmake -GNinja {} {}".format(args.vcvars,
-                                                         CMAKE_FLAGS[args.flavour],
-                                                         mxnet_root)
-            logging.info("Generating project with CMake:\n{}".format(cmd))
-            check_call(cmd, shell=True)
-
-            cmd = "\"{}\" && ninja".format(args.vcvars)
-            logging.info("Building:\n{}".format(cmd))
-
-            t0 = int(time.time())
-            ret = call(cmd, shell=True)
-
-            if ret != 0:
-                build_try += 1
-                logging.info("{} build(s) have failed".format(build_try))
-            else:
-                logging.info("Build flavour: {} complete in directory: \"{}\"".format(args.flavour, os.path.abspath(path)))
-                logging.info("Build took {}".format(datetime.timedelta(seconds=int(time.time() - t0))))
-                break
-    if ret == 0:
-        windows_package(args)
-    else:
+            tmpdirname = tempfile.mkdtemp()
+            os.chdir(tmpdirname)
+            r = requests.get('https://github.com/thrust/thrust/archive/1.9.8.zip', allow_redirects=True)
+            with open('thrust.zip', 'wb') as f:
+                f.write(r.content)
+            with zipfile.ZipFile('thrust.zip', 'r') as zip_ref:
+                zip_ref.extractall('.')
+            thrust_path = os.path.join(tmpdirname, "thrust-1.9.8")
+
+    with remember_cwd():
+        os.chdir(path)
+        env = os.environ.copy()
+        if 'GPU' in args.flavour:
+            env["CXXFLAGS"] = '/FS /MD /O2 /Ob2 /I {}'.format(thrust_path)
+        cmd = "\"{}\" && cmake -GNinja {} {}".format(args.vcvars,
+                                                     CMAKE_FLAGS[args.flavour],
+                                                     mxnet_root)
+        logging.info("Generating project with CMake:\n{}".format(cmd))
+        check_call(cmd, shell=True, env=env)
+
+        cmd = "\"{}\" && ninja".format(args.vcvars)
+        logging.info("Building:\n{}".format(cmd))
+
+        t0 = int(time.time())
+        ret = call(cmd, shell=True)
+
+    if ret != 0:
+        logging.info("Build failed")
         sys.exit(1)
 
+    logging.info("Build flavour: {} complete in directory: \"{}\"".format(args.flavour, os.path.abspath(path)))
+    logging.info("Build took {}".format(datetime.timedelta(seconds=int(time.time() - t0))))
+    windows_package(args)
+
 
 def windows_package(args):
     pkgfile = 'windows_package.7z'