diff --git a/README.md b/README.md index 29b5e7a16..b68866423 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,4 @@ -# Prophet: Automatic Forecasting Procedure - +# Prophet: Automatic Forecasting Procedure With python OPENCL/Multi-Core options ![Build](https://github.com/facebook/prophet/workflows/Build/badge.svg) [![PyPI Version](https://img.shields.io/pypi/v/prophet.svg)](https://pypi.python.org/pypi/prophet) diff --git a/python/setup.py b/python/setup.py index 20b7c69ad..772fb776c 100644 --- a/python/setup.py +++ b/python/setup.py @@ -4,11 +4,13 @@ # LICENSE file in the root directory of this source tree. import os +import sys import platform from pathlib import Path from shutil import copy, copytree, rmtree from typing import List - +import cmdstanpy.utils as utils +print(utils.cmdstan_path()) from setuptools import find_packages, setup, Extension from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py @@ -18,14 +20,25 @@ MODEL_DIR = "stan" MODEL_TARGET_DIR = os.path.join("prophet", "stan_model") - +MODEL_GPU = False +MODEL_MULTI_CORES = False CMDSTAN_VERSION = "2.26.1" BINARIES_DIR = "bin" BINARIES = ["diagnose", "print", "stanc", "stansummary"] TBB_PARENT = "stan/lib/stan_math/lib" TBB_DIRS = ["tbb", "tbb_2019_U8"] - - +if sys.platform.startswith('win'): + try: + OPEN_CL = "-L\"" + "%CUDA_PATH%/lib/Win32/OpenCL.lib\"" + except KeyError: + raise OSError("Environment variable is not set") +elif sys.platform.startswith('linux'): + try: + OPEN_CL = "-L/usr/local/cuda-11.8/lib64/ -lOpenCL" + except KeyError: + raise OSError("Environment variable is not set") +else: + raise OSError("Unsupported platform: {}".format(sys.platform)) def prune_cmdstan(cmdstan_dir: str) -> None: """ Keep only the cmdstan executables and tbb files (minimum required to run a cmdstanpy commands on a pre-compiled model). @@ -105,12 +118,40 @@ def build_cmdstan_model(target_dir): target_dir: Directory to copy the compiled model executable and core cmdstan files to. """ import cmdstanpy - + if MODEL_GPU and MODEL_MULTI_CORES: + cpp_options = { + 'STAN_THREADS': True, + 'STAN_CPP_OPTIMS': True, + 'STAN_OPENCL': True, + 'OPENCL_DEVICE_ID': 0, + 'OPENCL_PLATFORM_ID': 0, + "PRECOMPILED_HEADERS":False, + 'LDFLAGS': OPEN_CL+' -lOpenCL' + + } + elif MODEL_MULTI_CORES: + cpp_options = { + 'STAN_THREADS': True, + 'STAN_CPP_OPTIMS': True, + + } + elif MODEL_GPU: + cpp_options = { + 'STAN_CPP_OPTIMS': True, + 'STAN_OPENCL': True, + 'OPENCL_DEVICE_ID': 0, + 'OPENCL_PLATFORM_ID': 0, + "PRECOMPILED_HEADERS":False, + 'LDFLAGS': OPEN_CL+' -lOpenCL' + } cmdstan_dir = (Path(target_dir) / f"cmdstan-{CMDSTAN_VERSION}").resolve() install_cmdstan_deps(cmdstan_dir) model_name = "prophet.stan" target_name = "prophet_model.bin" - sm = cmdstanpy.CmdStanModel(stan_file=os.path.join(MODEL_DIR, model_name)) + if MODEL_GPU or MODEL_MULTI_CORES: + sm = cmdstanpy.CmdStanModel(stan_file=os.path.join(MODEL_DIR, model_name), cpp_options=cpp_options) + else: + sm = cmdstanpy.CmdStanModel(stan_file=os.path.join(MODEL_DIR, model_name)) copy(sm.exe_file, os.path.join(target_dir, target_name)) # Clean up