Windows support for autotvm - Do not merge

CMakeLists.txt

@@ -1,4 +1,9 @@
-cmake_minimum_required(VERSION 3.2)
+if(WIN32)
+  cmake_minimum_required(VERSION 3.9)
+else()
+  cmake_minimum_required(VERSION 3.2)
+endif()
+
 project(tvm C CXX)
 
 # Utility functions
@@ -442,6 +447,9 @@ endif(INSTALL_DEV)
 
 # More target definitions
 if(MSVC)
+  set_property(TARGET tvm PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
+  set_property(TARGET tvm_topi PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
+  set_property(TARGET tvm_runtime PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
   target_compile_definitions(tvm PRIVATE -DTVM_EXPORTS)
   target_compile_definitions(tvm_runtime PRIVATE -DTVM_EXPORTS)
 endif()

apps/cpp_rpc/rpc_env.cc

@@ -194,7 +194,7 @@ void WindowsShared(const std::string& output,
                    const std::string& options = "", 
                    const std::string& cc = "clang") {
   std::string cmd = cc;
-  cmd += " -O2 -flto=full -fuse-ld=lld-link -Wl,/EXPORT:__tvm_main__ -shared ";
+  cmd += " -O2 -flto=full -fuse-ld=lld-link -shared ";
   cmd += " -o " + output;
   for (const auto& file : files) {
     cmd += " " + file;

python/tvm/autotvm/measure/local_executor.py

@@ -17,8 +17,30 @@
 """Local based implementation of the executor using multiprocessing"""
 
 import signal
-
-from multiprocessing import Process, Queue
+import os
+
+if os.name == 'nt':
+    import queue as thread_queue
+    import threading
+    # Pathos uses dill, which can pickle things like functions
+    from pathos.helpers import ProcessPool
+    # On Windows, there is no fork(), a 'multiprocessing.Process'
+    # or ProcessPool has to 'build up' the script from scratch, we
+    # set these environment variables so each python.exe process
+    # does not allocate unneeded threads
+    os.environ['OMP_NUM_THREADS'] = "1"
+    os.environ['TVM_NUM_THREADS'] = "1"
+    # numpy seems to honor this
+    os.environ['MKL_NUM_THREADS'] = "1"
+
+    # Since there is no fork() on Windows, to mitigate performance impact
+    # we will use a process pool for executers, vs the *nix based systems
+    # that will fork() a new process for each executor
+    EXECUTOR_POOL = None
+
+#pylint: disable=wrong-import-position
+#pylint: disable=ungrouped-imports
+from multiprocessing import Process, Queue, cpu_count
 try:
     from queue import Empty
 except ImportError:
@@ -30,7 +52,8 @@
     psutil = None
 
 from . import executor
-
+#pylint: enable=ungrouped-imports
+#pylint: enable=wrong-import-position
 
 def kill_child_processes(parent_pid, sig=signal.SIGTERM):
     """kill all child processes recursively"""
@@ -68,6 +91,26 @@ def call_with_timeout(queue, timeout, func, args, kwargs):
     p.terminate()
     p.join()
 
+if os.name == 'nt':
+    def call_from_pool(func, args, kwargs, timeout, env): # pylint: disable=unused-argument
+        """A wrapper to support timeout of a function call for a pool process"""
+
+        # Restore environment variables from parent
+        for key, val in env.items():
+            os.environ[key] = val
+
+        queue = thread_queue.Queue(2)
+
+        # We use a thread here for Windows, because starting up a new Process can be heavy
+        # This isn't as clean as the *nix implementation, which can kill a process that
+        # has timed out
+        thread = threading.Thread(target=_execute_func, args=(func, queue, args, kwargs))
+        thread.start()
+        thread.join()
+        queue.put(executor.TimeoutError())
+
+        res = queue.get()
+        return res
 
 class LocalFuture(executor.Future):
     """Local wrapper for the future
@@ -119,6 +162,31 @@ def done(self):
     def get(self, timeout=None):
         return self._result
 
+if os.name == 'nt':
+    class LocalFuturePool(executor.Future):
+        """Local wrapper for the future using a Process pool
+
+        Parameters
+        ----------
+        thread: threading.Thread
+            Thread for running this task
+        pool_results: result from Pool.apply_async
+            queue for receiving the result of this task
+        """
+        def __init__(self, pool_results):
+            self._done = False
+            self._pool_results = pool_results
+
+        def done(self):
+            return self._done
+
+        def get(self, timeout=None):
+            try:
+                res = self._pool_results.get()
+            except Empty:
+                raise executor.TimeoutError()
+            self._done = True
+            return res
 
 class LocalExecutor(executor.Executor):
     """Local executor that runs workers on the same machine with multiprocessing.
@@ -145,8 +213,25 @@ def submit(self, func, *args, **kwargs):
         if not self.do_fork:
             return LocalFutureNoFork(func(*args, **kwargs))
 
-        queue = Queue(2)
-        process = Process(target=call_with_timeout,
-                          args=(queue, self.timeout, func, args, kwargs))
-        process.start()
-        return LocalFuture(process, queue)
+        if os.name != 'nt':
+            queue = Queue(2)
+            process = Process(target=call_with_timeout,
+                              args=(queue, self.timeout, func, args, kwargs))
+            process.start()
+            return LocalFuture(process, queue)
+
+        global EXECUTOR_POOL
+
+        if EXECUTOR_POOL is None:
+            # We use a static pool for executor processes because Process.start(entry)
+            # is so slow on Windows, we lose a lot of parallelism.
+            # Right now cpu_count() is used, which isn't optimal from a user configuration
+            # perspective, but is reasonable at this time.
+            EXECUTOR_POOL = ProcessPool(cpu_count() * 2)
+
+        # Windows seemed to be missing some valuable environ variables
+        # on the pool's process side.  We might be able to get away with
+        # just sending the PATH variable, but for now, we just clone our env
+        return LocalFuturePool(EXECUTOR_POOL.apply_async(call_from_pool,
+                                                         (func, args, kwargs,
+                                                          self.timeout, os.environ.copy())))

python/tvm/autotvm/measure/measure_methods.py

@@ -27,7 +27,9 @@
 import os
 import threading
 import time
-from random import getrandbits
+# Import random directly because it appears dill will pull in 'getrandbits'
+# and it will always get the same random number. Using random.getrandbits fixes
+import random
 from collections import namedtuple
 import tempfile
 
@@ -331,10 +333,12 @@ def set_task(self, task):
         from ...rpc.tracker import Tracker
         from ...rpc.server import Server
 
+        # Windows will not let you connect to 0.0.0.0
+        local_address = '0.0.0.0' if os.name != 'nt' else '127.0.0.1'
         self.task = task
-        tracker = Tracker('0.0.0.0', port=9000, port_end=10000, silent=True)
+        tracker = Tracker(local_address, port=9000, port_end=10000, silent=True)
         device_key = '$local$device$%d' % tracker.port
-        server = Server('0.0.0.0', port=9000, port_end=10000,
+        server = Server(local_address, port=9000, port_end=10000,
                         key=device_key,
                         use_popen=True, silent=True,
                         tracker_addr=(tracker.host, tracker.port))
@@ -392,6 +396,7 @@ def _wrap_build_func(build_func):
         raise AttributeError("Expect build_func to have the attribute output_format.")
     output_format = build_func.output_format
 
+
     def _wrapped(measure_input, tmp_dir, **kwargs):
         """
         Wrapped build func.
@@ -407,7 +412,7 @@ def _wrapped(measure_input, tmp_dir, **kwargs):
         tic = time.time()
         try:
             filename = os.path.join(tmp_dir, "tmp_func_%0x.%s" % (
-                getrandbits(64), output_format))
+                random.getrandbits(64), output_format))
             # TODO(tvm-team) consider linline _build_func_common
             func, arg_info = _build_func_common(measure_input, **kwargs)
             func.export_library(filename, build_func)

python/tvm/autotvm/task/relay_integration.py

@@ -135,8 +135,12 @@ def extract_from_multiple_program(mods, params, target, target_host=None, ops=No
             # wrap build call in thread to avoid multiprocessing problems
             build_thread = threading.Thread(target=_lower,
                                             args=(mod, target, param))
+            # Stack would overflow on some platforms (Windows) on some models
+            old_stack_size = threading.stack_size(1024 * 1024 * 3)
             build_thread.start()
             build_thread.join()
+            # Restore stacksize to original
+            threading.stack_size(old_stack_size)
 
         logger.disabled = old_state
 

python/tvm/autotvm/task/task.py

@@ -21,6 +21,8 @@
 func is a state-less function, or a string that
 registers the standard task.
 """
+
+import os
 import numpy as np
 
 from tvm import target as _target
@@ -169,7 +171,9 @@ def __getstate__(self):
             "config_space": self.config_space,
             "flop": self.flop,
             "target": self.target,
-            "target_host": self.target_host
+            "target_host": self.target_host,
+            # On Windows we will use, dill, which can pickle functions
+            "func": self.func if os.name == 'nt' else None
         }
 
     def __setstate__(self, state):

python/tvm/autotvm/tuner/xgboost_cost_model.py

@@ -20,7 +20,15 @@
 import multiprocessing
 import logging
 import time
-
+import os
+
+if os.name == 'nt':
+    # Pathos' Pool does pickling via dill, which can pickle
+    # functions, which is required because Windows doesn't
+    # support fork()
+    from pathos.helpers import mp as pathos_multiprocess
+    from pathos.helpers import ProcessPool
+#pylint: disable=wrong-import-position
 import numpy as np
 try:
     import xgboost as xgb
@@ -31,7 +39,7 @@
 from ..util import get_rank
 from .metric import max_curve, recall_curve, cover_curve
 from .model_based_tuner import CostModel, FeatureCache
-
+#pylint: enable=wrong-import-position
 logger = logging.getLogger('autotvm')
 
 class XGBoostCostModel(CostModel):
@@ -153,14 +161,60 @@ def _reset_pool(self, space, target, task):
 
         self._close_pool()
 
-        # use global variable to pass common arguments
-        global _extract_space, _extract_target, _extract_task
-        _extract_space = space
-        _extract_target = target
-        _extract_task = task
-        self.pool = multiprocessing.Pool(self.num_threads)
+        if os.name == 'nt':
+            # For Windows, we need space, target, task to be pickled and set on the
+            # Pool's process side, where the *nix impl simply sets globals
+            # then forks.
+            # To ensure each process in the pool is properly set, we have to do
+            # some synchronization by sending an async call, setting space, target and task, then
+            # waiting for the queue to have an item set
+            num_threads = self.num_threads
+            pool_size = num_threads if num_threads is not None else multiprocessing.cpu_count()
+            if self.pool is None:
+                self.pool = ProcessPool(pool_size)
+            manager = pathos_multiprocess.Manager()
+            pipe_syncs = []
+
+            # A simple pathos.map would be cleaner, but it seems that in some cases,
+            # some of the pools processes will be missed, with some processes running
+            # the method twice.  It seems that just passing a Queue in this manner,
+            # hits all the processes in the pool. Some assertion could be built to verify
+            for _ in range(pool_size):
+                queue = manager.Queue(1)
+                results = {
+                    "queue": queue,
+                    "apipe": self.pool.apply_async(_set_pool_process_state,
+                                                   args=(space, target, task, queue))
+                }
+                pipe_syncs.append(results)
+
+            # wait loop until all async calls have completed
+            while True:
+                all_ready = True
+                for pipe_sync in pipe_syncs:
+                    if pipe_sync["apipe"].ready() is False:
+                        all_ready = False
+                        break
+                if all_ready:
+                    break
+                time.sleep(0.05)
+            # complete the async requests on the pool
+            for pipe_sync in pipe_syncs:
+                pipe_sync["apipe"].get()
+                # This may not be needed
+                pipe_sync["queue"].get(block=True)
+        else:
+            # use global variable to pass common arguments
+            global _extract_space, _extract_target, _extract_task
+            _extract_space = space
+            _extract_target = target
+            _extract_task = task
+            self.pool = multiprocessing.Pool(self.num_threads)
+
+    def _close_pool(self, force_close=False):
+        if os.name == 'nt' and not force_close:
+            return
 
-    def _close_pool(self):
         if self.pool:
             self.pool.terminate()
             self.pool.join()
@@ -195,7 +249,6 @@ def fit(self, xs, ys, plan_size):
                 self.base_model = None
             else:
                 dtrain.set_base_margin(discount * self.base_model.predict(xs, output_margin=True))
-
         self.bst = xgb.train(self.xgb_params, dtrain,
                              num_boost_round=8000,
                              callbacks=[custom_callback(
@@ -324,13 +377,23 @@ def _get_feature(self, indexes):
         return ret
 
     def __del__(self):
-        self._close_pool()
+        self._close_pool(force_close=True)
 
 
 _extract_space = None
 _extract_target = None
 _extract_task = None
 
+if os.name == 'nt':
+    def _set_pool_process_state(space, target, task, sync_queue):
+        """sets process state for when fork() is not available """
+        global _extract_space, _extract_target, _extract_task
+        _extract_space = space
+        _extract_target = target
+        _extract_task = task
+        # Notify caller that we are done. We may be able to remove this
+        sync_queue.put(None)
+
 def _extract_itervar_feature_index(index):
     """extract iteration var feature for an index in extract_space"""
     try:
@@ -423,7 +486,7 @@ def custom_callback(stopping_rounds, metric, fevals, evals=(), log_file=None,
     from xgboost.core import EarlyStopException
     from xgboost.callback import _fmt_metric
     from xgboost.training import aggcv
-
+    #pylint: enable=import-outside-toplevel
     state = {}
     metric_shortname = metric.split("-")[1]
 
@@ -472,7 +535,6 @@ def callback(env):
                 res = [x.split(':') for x in bst_eval.split()]
                 for kv in res[1:]:
                     res_dict[kv[0]] = [float(kv[1])]
-
         eval_res = []
         keys = list(res_dict.keys())
         keys.sort(key=lambda x: x if metric_shortname not in x else "a" + x)

python/tvm/contrib/cc.py

@@ -211,6 +211,7 @@ def _windows_shared(output, objects, options):
     except FileNotFoundError:
         raise RuntimeError("Can not find cl.exe,"
                            "please run this in Vistual Studio Command Prompt.")
+    print(py_str(out))
     if proc.returncode != 0:
         msg = "Compilation error:\n"
         msg += py_str(out)
@@ -226,7 +227,6 @@ def _windows_shared(output, objects, options):
         if obj.endswith(".o"):
             link_cmd += [obj]
 
-    link_cmd += ["-EXPORT:__tvm_main__"]
     link_cmd += [temp_path + "dllmain.obj"]
     link_cmd += ["-out:" + output]