Build packages for npcomp-torch.

* Adds a minimal setup.py for frontends/pytorch
* Makes npcomp-core export its headers and libraries
* Adds a script to build packages.
* Adds CI step to package and smoke test.
* Will need some more tweaks and coordination prior to deploying (version locking etc).

.github/workflows/buildAndTest.yml

@@ -7,16 +7,20 @@ jobs:
     name: Build and Test (Release Asserts)
     runs-on: ubuntu-20.04
     steps:
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.9
     - name: Get npcomp
       uses: actions/checkout@v2
       with:
         submodules: 'true'
     - name: Install python depends
       run: |
-        python3 -m pip install -r $GITHUB_WORKSPACE/external/llvm-project/mlir/python/requirements.txt
+        python -m pip install -r $GITHUB_WORKSPACE/external/llvm-project/mlir/python/requirements.txt
     - name: Install pytorch_nightly depends
       run: |
-        python3 -m pip install --pre 'torch==1.10.0.dev20210630+cpu' 'torchvision==0.11.0.dev20210630+cpu' -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+        python -m pip install --pre 'torch==1.10.0.dev20210630+cpu' 'torchvision==0.11.0.dev20210630+cpu' -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
     - name: Install Ninja
       uses: llvm/actions/install-ninja@55d844821959226fab4911f96f37071c1d4c3268
     - name: Get Submodule Hash
@@ -26,7 +30,7 @@ jobs:
     - name: Ccache for C++ compilation
       uses: hendrikmuhs/ccache-action@4687d037e4d7cf725512d9b819137a3af34d39b3
       with:
-        key: ${{ runner.os }}-${{ steps.get-submodule-hash.outputs.hash }}
+        key: ${{ runner.os }}-clangreleaseasserts-${{ steps.get-submodule-hash.outputs.hash }}
     - name: Build and Test npcomp (Assert)
       run: |
         cd $GITHUB_WORKSPACE
@@ -37,9 +41,44 @@ jobs:
           -DCMAKE_LINKER=lld \
           -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
           -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
-          -DPython3_EXECUTABLE=/usr/bin/python3 \
+          -DPython3_EXECUTABLE=$(which python) \
           -DLLVM_ENABLE_ASSERTIONS=ON \
           -DLLVM_TARGETS_TO_BUILD=host \
           -DNPCOMP_ENABLE_PYTORCH=ON
         ninja
         ninja check-npcomp check-frontends-pytorch
+
+  build_wheels:
+    name: Build Wheels and Smoketest
+    runs-on: ubuntu-20.04
+    steps:
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.9
+    - name: Get npcomp
+      uses: actions/checkout@v2
+      with:
+        submodules: 'true'
+    - name: Install python dependencies
+      run: |
+        python -m pip install --pre 'torch==1.10.0.dev20210630+cpu' 'torchvision==0.11.0.dev20210630+cpu' -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+        python -m pip install wheel
+    - name: Install Ninja
+      uses: llvm/actions/install-ninja@55d844821959226fab4911f96f37071c1d4c3268
+    - name: Get Submodule Hash
+      id: get-submodule-hash
+      run: echo "::set-output name=hash::$(md5sum $(git submodule status))"
+      shell: bash
+    - name: Ccache for C++ compilation
+      uses: hendrikmuhs/ccache-action@4687d037e4d7cf725512d9b819137a3af34d39b3
+      with:
+        key: ${{ runner.os }}-wheel-${{ steps.get-submodule-hash.outputs.hash }}
+    - name: Build wheels
+      run: |
+        cd $GITHUB_WORKSPACE
+        ./build_tools/build_python_wheels.sh
+    - uses: actions/upload-artifact@v2
+      with:
+        name: ubuntu-20.04-wheels
+        path: wheelhouse/

.gitignore

@@ -3,10 +3,7 @@
 .env
 *.code-workspace
 
-/build
-/build_*
-build-mlir
-install-mlir
+/build/
 __pycache__
 
 .pytype

README.md

@@ -9,6 +9,20 @@ The NPComp project aims to provide tooling for compiling numerical python progra
 
 In addition to providing a bridge to a variety of Python based numerical programming frameworks, NPComp also directly develops components for tracing and compilation of generic Python program fragments.
 
+## Quick Build and Install with PyTorch
+
+Instructions below go in to detail about how to get a functioning development
+setup. But to just build and install into a local python instance, the
+instructions are simple. This support is new and we still need to get
+dependencies pinned and make these packages distributable. This should work
+locally, though:
+
+```
+python -m pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+
+./build_tools/build_python_wheels.sh
+```
+
 ## Project Communication
 
 - `#mlir-npcomp` channel on the LLVM [Discord](https://discord.gg/xS7Z362)

build_tools/build_python_wheels.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+set -e
+
+if [ -z "$PYTHON" ]; then
+  PYTHON="$(which python)"
+fi
+version="$("$PYTHON" --version)"
+echo "Using python: $PYTHON (version $version)"
+
+repo_root="$(cd $(dirname $0)/.. && pwd)"
+wheelhouse="$repo_root/wheelhouse"
+mkdir -p $wheelhouse
+cd $wheelhouse
+
+echo "---- BUILDING npcomp-core ----"
+CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache \
+$PYTHON -m pip wheel "${repo_root}" \
+  --use-feature=in-tree-build \
+  -w "$wheelhouse" -v
+
+echo "---- INSTALLING npcomp-core ----"
+$PYTHON -m pip install -f "$wheelhouse" --force-reinstall npcomp-core
+
+echo "---- BUILDING npcomp-torch ----"
+$PYTHON -m pip wheel "${repo_root}/frontends/pytorch" \
+  --use-feature=in-tree-build \
+  -w "$wheelhouse" -v
+
+echo "---- INSTALLING npcomp-torch ----"
+$PYTHON -m pip install -f "$wheelhouse" --force-reinstall npcomp-torch
+
+echo "---- QUICK SMOKE TEST ----"
+$PYTHON $repo_root/frontends/pytorch/test/torchscript_e2e_test/basic.py

external/llvm-project

@@ -1 +1 @@
-Subproject commit 5b2e7f50a6798fd9b9c79d9d62fdebcd9e78525b
+Subproject commit cf36ab1d6c39e80a70b5a1dc80120cccccb5301c

frontends/pytorch/.gitignore

@@ -0,0 +1,2 @@
+/build
+/dist

frontends/pytorch/python/torch_mlir/__init__.py

@@ -6,6 +6,11 @@
 # and binds names locally. It exists to allow for customization of behavior
 # prior to loading shared objects.
 
+import torch
+import npcomp
+
+# Our native extension is not self-contained. It references libraries which
+# must come in via the above first.
 from _torch_mlir import *
 
 

frontends/pytorch/setup.py

@@ -0,0 +1,45 @@
+# Build/install the npcomp-torch package.
+# This uses PyTorch's setuptools support and requires an existing installation
+# of npcomp-core in order to access its headers/libraries.
+
+from pathlib import Path
+
+from setuptools import find_packages, setup, Extension
+from torch.utils import cpp_extension
+
+try:
+  from npcomp import build as npcomp_build
+except ModuleNotFoundError:
+  raise ModuleNotFoundError(
+    f"Could not import npcomp.build "
+    f"(do you have the npcomp-core package installed)")
+
+# Get our sources.
+this_dir = Path(__file__).parent
+extension_sources = [str(p) for p in this_dir.joinpath("csrc").rglob("*.cpp")]
+
+# Npcomp bits.
+include_dirs = npcomp_build.get_include_dirs()
+lib_dirs = npcomp_build.get_lib_dirs()
+npcomp_libs = [npcomp_build.get_capi_link_library_name()]
+
+setup(
+  name="npcomp-torch",
+  ext_modules=[
+    cpp_extension.CppExtension(
+      name="_torch_mlir",
+      sources=extension_sources,
+      include_dirs=include_dirs,
+      library_dirs=lib_dirs,
+      libraries=npcomp_libs),
+  ],
+  cmdclass={
+    "build_ext": cpp_extension.BuildExtension
+  },
+  package_dir={
+    "": "./python",
+  },
+  packages=find_packages("./python", include=[
+    "torch_mlir", "torch_mlir.*",
+  ]),
+)

python/CMakeLists.txt

@@ -24,15 +24,14 @@ add_dependencies(NPCOMPPythonResources
 # Declare sources
 ################################################################################
 
-set(_addl_extension_sources)
-if(NPCOMP_ENABLE_REFJIT)
-  list(APPEND _addl_extension_sources "${CMAKE_CURRENT_SOURCE_DIR}/RefJITBackend.cpp")
-endif()
+declare_mlir_python_sources(NPCOMPPythonSources)
 
-declare_mlir_python_sources(NPCOMPPythonSources
+declare_mlir_python_sources(NPCOMPPythonSources.Core
+  ADD_TO_PARENT NPCOMPPythonSources
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
   SOURCES
     npcomp/__init__.py
+    npcomp/build.py
     npcomp/decorators.py
     npcomp/exporter.py
     npcomp/smoketest.py
@@ -45,11 +44,27 @@ declare_mlir_python_sources(NPCOMPPythonSources
     npcomp/tracing/*.py
     npcomp/utils/*.py
 )
+
 declare_mlir_python_sources(NPCOMPPythonSources.Dialects
   ADD_TO_PARENT NPCOMPPythonSources
 )
 declare_mlir_python_sources(NPCOMPPythonExtensions)
 
+declare_mlir_python_sources(NPCOMPPythonCAPIHeaderSources
+  ROOT_DIR "${MLIR_NPCOMP_SOURCE_DIR}/include"
+  SOURCES_GLOB "npcomp-c/*.h"
+  DEST_PREFIX "_mlir_libs/include"
+)
+
+################################################################################
+# Extensions
+################################################################################
+
+set(_addl_extension_sources)
+if(NPCOMP_ENABLE_REFJIT)
+  list(APPEND _addl_extension_sources "${CMAKE_CURRENT_SOURCE_DIR}/RefJITBackend.cpp")
+endif()
+
 declare_mlir_python_extension(NPCOMPPythonExtensions.Core
   MODULE_NAME _npcomp
   ADD_TO_PARENT NPCOMPPythonExtensions
@@ -126,7 +141,9 @@ add_mlir_python_modules(NPCOMPMLIRPythonModules
     MLIRPythonExtension.AllPassesRegistration
     # We need the npcomp extensions co-located with the MLIR extensions. When
     # the namespace is unified, this moves to the below.
+    MLIRPythonCAPIHeaderSources
     NPCOMPPythonExtensions
+    NPCOMPPythonCAPIHeaderSources
   COMMON_CAPI_LINK_LIBS
     NPCOMPPythonCAPI
 )

python/npcomp/build.py

@@ -0,0 +1,50 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""Exports configuration for the package and settings for building libraries."""
+
+import os
+import platform
+
+from mlir._mlir_libs import get_include_dirs, get_lib_dirs
+
+__all__ = [
+  "get_include_dirs",
+  "get_lib_dirs",
+]
+
+
+def get_capi_link_library_name() -> str:
+  """Gets the library name of the CAPI shared library to link against."""
+  return "NPCOMPPythonCAPI"
+
+
+def get_capi_link_library_path() -> str:
+  """Returns an absolute path to the CAPI shared library.
+
+  This should be preferred when seeking to create a non relocatable package
+  as it eliminates the possibility of interference of similar named libraries
+  on the link path.
+
+  Raises:
+    ValueError: If the library cannot be found.
+  """
+  system = platform.system()
+  lib_prefix = "lib"
+  lib_suffix = ".so"
+  if system == "Darwin":
+    lib_suffix = ".dylib"
+  elif system == "Windows":
+    lib_prefix = ""
+    lib_suffix = ".lib"
+  lib_filename = f"{lib_prefix}{get_capi_link_library_name()}{lib_suffix}"
+
+  for lib_dir in get_lib_dirs():
+    full_path = os.path.join(lib_dir, lib_filename)
+    if os.path.exists(full_path): return full_path
+
+  raise ValueError(
+    f"Unable to find npcomp development library {lib_filename} in "
+    f"{get_lib_dirs()}")
+

setup.py

@@ -77,6 +77,7 @@ setup(
     author_email="stellaraccident@gmail.com",
     description="NPComp Core",
     long_description="",
+    include_package_data=True,
     ext_modules=[
         CMakeExtension("mlir._mlir_libs._mlir"),
         CMakeExtension("mlir._mlir_libs._npcomp"),

test/development_files.py

@@ -0,0 +1,11 @@
+# RUN: %PYTHON %s 2>&1
+
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from npcomp import build
+
+assert build.get_include_dirs()
+assert build.get_lib_dirs()
+print("CAPI Path:", build.get_capi_link_library_path())