Remove aten submodule, add pytorch submodule and get aten from there. Also make TC work with the latest aten AP and switch python API to use pytorch-nightly. CI is switched from gcc 4.8.* to gcc 4.9 to maintain compatibility with the pytorch binary which is using gcc 4.9

Author: prigoyal

.circleci/config.yml

@@ -4,7 +4,7 @@ jobs:
     working_directory: ~/TensorComprehensions
     resource_class: xlarge
     docker:
-      - image: tensorcomprehensions/linux-trusty-gcc4.8-cuda8-cudnn6-py3-conda:x86
+      - image: tensorcomprehensions/linux-trusty-gcc4.9-cuda8-cudnn7-py3-conda:1
 
     steps:
       - checkout
@@ -26,7 +26,7 @@ jobs:
 
       - restore_cache:
           keys:
-            - v1-aten-{{ checksum ".git/modules/third-party/ATen/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1404
+            - v1-pytorch-{{ checksum ".git/modules/third-party/pytorch/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1404
 
       - restore_cache:
           keys:
@@ -68,7 +68,7 @@ jobs:
             - third-party/caffe2/build_host_protoc
 
       - save_cache:
-          key: v1-aten-{{ checksum ".git/modules/third-party/ATen/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1404
+          key: v1-aten-{{ checksum ".git/modules/third-party/pytorch/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1404
           paths:
             - third-party-install/share/ATen
             - third-party-install/include/ATen
@@ -78,19 +78,9 @@ jobs:
             - third-party-install/include/THCUNN
             - third-party-install/include/THNN
             - third-party-install/include/THS
+            - third-party-install/include/cpuinfo.h
             - third-party-install/lib/libATen.so
-            - third-party-install/lib/libATen.so.1
-            - third-party-install/lib/libTH.so
-            - third-party-install/lib/libTH.so.0
-            - third-party-install/lib/libTHC.so
-            - third-party-install/lib/libTHC.so.0
-            - third-party-install/lib/libTHCS.so
-            - third-party-install/lib/libTHCS.so.1
-            - third-party-install/lib/libTHCUNN.so
-            - third-party-install/lib/libTHNN.so
-            - third-party-install/lib/libTHS.so
-            - third-party-install/lib/libTHS.so.1
-            - third-party/ATen/build/src/ATen/test
+            - third-party/pytorch/aten/build/src/ATen/test/
 
       - save_cache:
           key: v1-isl-{{ checksum ".git/modules/third-party/islpp/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1404
@@ -125,7 +115,7 @@ jobs:
     working_directory: ~/TensorComprehensions
     resource_class: xlarge
     docker:
-      - image: tensorcomprehensions/linux-xenial-gcc5-cuda9-cudnn7-py3:x86
+      - image: tensorcomprehensions/linux-xenial-gcc5-cuda9-cudnn7-py3:1
 
     steps:
       - checkout
@@ -188,7 +178,7 @@ jobs:
             - third-party-install/lib/libcaffe2_gpu.so
 
       - save_cache:
-          key: v1-aten-{{ checksum ".git/modules/third-party/ATen/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1604
+          key: v1-aten-{{ checksum ".git/modules/third-party/pytorch/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1604
           paths:
             - third-party-install/share/ATen
             - third-party-install/include/ATen
@@ -198,19 +188,9 @@ jobs:
             - third-party-install/include/THCUNN
             - third-party-install/include/THNN
             - third-party-install/include/THS
+            - third-party-install/include/cpuinfo.h
             - third-party-install/lib/libATen.so
-            - third-party-install/lib/libATen.so.1
-            - third-party-install/lib/libTH.so
-            - third-party-install/lib/libTH.so.0
-            - third-party-install/lib/libTHC.so
-            - third-party-install/lib/libTHC.so.0
-            - third-party-install/lib/libTHCS.so
-            - third-party-install/lib/libTHCS.so.1
-            - third-party-install/lib/libTHCUNN.so
-            - third-party-install/lib/libTHNN.so
-            - third-party-install/lib/libTHS.so
-            - third-party-install/lib/libTHS.so.1
-            - third-party/ATen/build/src/ATen/test
+            - third-party/pytorch/aten/build/src/ATen/test/
 
       - save_cache:
           key: v1-isl-{{ checksum ".git/modules/third-party/islpp/HEAD" }}-{{ checksum "build.sh" }}-{{ checksum ".circleci/config.yml" }}-{{ arch }}-1604

.gitmodules

@@ -23,9 +23,6 @@
 [submodule "third-party/glog"]
 	path = third-party/glog
 	url = https://github.com/google/glog.git
-[submodule "third-party/ATen"]
-	path = third-party/ATen
-	url = https://github.com/zdevito/aten.git
 [submodule "third-party/dlpack"]
 	path = third-party/dlpack
 	url = https://github.com/dmlc/dlpack.git
@@ -35,3 +32,6 @@
 [submodule "third-party/pybind11"]
  	path = third-party/pybind11
 	url = https://github.com/pybind/pybind11.git
+[submodule "third-party/pytorch"]
+	path = third-party/pytorch
+	url = https://github.com/pytorch/pytorch.git

.jenkins/build.sh

@@ -70,9 +70,9 @@ if [[ "$DISTRIB_RELEASE" == 14.04 ]]; then
     echo "Building TC in conda env"
     conda create -y --name tc-env python=3.6
     source activate tc-env
-    conda install -y -c pytorch pytorch
-    conda install -y pyyaml
+    conda install -y pyyaml mkl-include
     conda install -yc conda-forge pytest
+    conda install -y pytorch-nightly=2018.04.17 -c pytorch
     WITH_PYTHON_C2=OFF CORES=$(nproc) CLANG_PREFIX=/usr/local/clang+llvm-tapir5.0 BUILD_TYPE=Release ./build.sh --all
   else
     echo "Building TC in non-conda env"
@@ -85,9 +85,9 @@ if [[ "$DISTRIB_RELEASE" == 16.04 ]]; then
     echo "Building TC in conda env"
     conda create -y --name tc-env python=3.6
     source activate tc-env
-    conda install -y pytorch cuda90 -c pytorch
-    conda install -y pyyaml
+    conda install -y pyyaml mkl-include
     conda install -yc conda-forge pytest
+    conda install -y pytorch-nightly=2018.04.17 cuda90 -c pytorch
     WITH_PYTHON_C2=OFF CORES=$(nproc) CLANG_PREFIX=/usr/local/clang+llvm-tapir5.0 BUILD_TYPE=Release ./build.sh --all
   else
     echo "Building TC in non-conda env"

.jenkins/run_test.sh

@@ -8,27 +8,12 @@ source /etc/lsb-release
 # condition: if 16.04 and conda, run only python tests
 # condition: if any and non-conda, run test.sh only
 
-# TODO: modify 2LUT tests from example_MLP_model and enable on CI
-
-if [[ "$DISTRIB_RELEASE" == 14.04 ]]; then
+if [[ $(conda --version | wc -c) -ne 0 ]]; then
+  echo "Running TC PyTorch tests"
+  ./test_python/run_test.sh
+else
+  # TODO: modify 2LUT tests from example_MLP_model and enable on CI
   echo "Running TC backend tests"
   FILTER_OUT=MLP_model ./test.sh
   ./build/tc/benchmarks/MLP_model --gtest_filter=-*2LUT*
-  if [[ $(conda --version | wc -c) -ne 0 ]]; then
-    source activate tc-env
-    echo "Running TC PyTorch tests"
-    ./test_python/run_test.sh
-  fi
-fi
-
-if [[ "$DISTRIB_RELEASE" == 16.04 ]]; then
-  if [[ $(conda --version | wc -c) -ne 0 ]]; then
-    echo "Running TC PyTorch tests"
-    source activate tc-env
-    ./test_python/run_test.sh
-  else
-    echo "Running TC backend tests"
-    FILTER_OUT=MLP_model ./test.sh
-    ./build/tc/benchmarks/MLP_model --gtest_filter=-*2LUT*
-  fi
 fi

CMakeLists.txt

@@ -203,12 +203,12 @@ execute_process(COMMAND "python" "-c"
     OUTPUT_VARIABLE PYTHON_SITE_PACKAGES OUTPUT_STRIP_TRAILING_WHITESPACE)
 message(STATUS "PYTHON site packages: \n${PYTHON_SITE_PACKAGES}")
 
-# if PyTorch is installed, we get libATen.so.1 from there, otherwise build it
+# if PyTorch is installed, we get libATen.so from there, otherwise build it
 if (__torch_install EQUAL 0)
   message(STATUS "TORCH INSTALLED, linking to ATen from PyTorch")
   set(ATEN_INCLUDE_DIR "${PYTHON_SITE_PACKAGES}/torch/lib/include")
   include_directories(${ATEN_INCLUDE_DIR})
-  find_library(ATEN_LIBRARIES NAMES libATen.so.1 PATHS ${PYTHON_SITE_PACKAGES}/torch/lib)
+  find_library(ATEN_LIBRARIES NAMES libATen.so PATHS ${PYTHON_SITE_PACKAGES}/torch/lib)
   link_directories("${PYTHON_SITE_PACKAGES}/../../")
   message(STATUS "Found ATen.so file: ${ATEN_LIBRARIES}")
 else()

build.sh

@@ -220,8 +220,8 @@ function install_glog() {
 }
 
 function install_aten() {
-  mkdir -p ${TC_DIR}/third-party/ATen/build || exit 1
-  cd       ${TC_DIR}/third-party/ATen/build || exit 1
+  mkdir -p ${TC_DIR}/third-party/pytorch/aten/build || exit 1
+  cd       ${TC_DIR}/third-party/pytorch/aten/build || exit 1
 
   if ! test ${USE_CONTBUILD_CACHE} || [ ! -d "${INSTALL_PREFIX}/include/ATen" ]; then
 
@@ -234,18 +234,18 @@ function install_aten() {
     fi
 
     # ATen errors out when using many threads for building - use 1 core for now
-    if should_rebuild ${TC_DIR}/third-party/ATen ${TC_DIR}/third-party/.aten_build_cache; then
+    if should_rebuild ${TC_DIR}/third-party/pytorch/aten ${TC_DIR}/third-party/.aten_build_cache; then
       echo "Installing ATen"
 
       if should_reconfigure .. .build_cache; then
         echo "Reconfiguring ATen"
         export PYTORCH_PYTHON=${PYTHON}
-        ${CMAKE_VERSION} ../aten -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DHAS_C11_ATOMICS=OFF -DNO_CUDA=${ATEN_NO_CUDA}
+        ${CMAKE_VERSION} .. -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DHAS_C11_ATOMICS=OFF -DNO_CUDA=${ATEN_NO_CUDA}
       fi
       make -j $CORES -s || exit 1
 
       set_cache .. .build_cache
-      set_bcache ${TC_DIR}/third-party/ATen ${TC_DIR}/third-party/.aten_build_cache
+      set_bcache ${TC_DIR}/third-party/pytorch/aten ${TC_DIR}/third-party/.aten_build_cache
     fi
 
     make install -j $CORES -s || exit 1
@@ -369,6 +369,7 @@ function install_tc_python() {
     else
       if [[ $(conda --version | wc -c) -ne 0 ]]; then
           echo "Found conda, going to install Python packages"
+          conda install -y mkl-include
           cd ${TC_DIR}
           export CONDA_PYTHON=$(which python3)
           echo "CONDA_PYTHON: ${CONDA_PYTHON}"

docker/README.md

@@ -1,20 +1,20 @@
 # How to build Tensor Comprehensions docker images
 
-Tensor Comprehensions supports gcc4.8, gcc5, cuda8-cudnn6, cuda9-cudnn7, Xenial, Trusty
+Tensor Comprehensions supports gcc4.9, gcc5, cuda8-cudnn6, cuda9-cudnn7, Xenial, Trusty
 and building in Conda/ Non-conda environment. Below are the steps you can follow
 to build the images for various combinations. This directory contains everything needed to build the Docker images that are used in our CI as well. The `Dockerfile`s are parameterized to build the image based on the build environment. The different configurations are identified by a freeform string that we call a build environment. This string is persisted in each image as the BUILD_ENVIRONMENT environment variable.
 
 ## Supported Build environments
 
 Non-conda env and no python installs:
 
-`linux-trusty-gcc4.8-cuda8-cudnn7-py3`
+`linux-trusty-gcc4.9-cuda8-cudnn7-py3`
 
 `linux-xenial-gcc5-cuda9-cudnn7-py3`
 
 Conda env with PyTorch conda install:
 
-`linux-trusty-gcc4.8-cuda8-cudnn7-py3-conda` [Run both C++ and Python tests]
+`linux-trusty-gcc4.9-cuda8-cudnn7-py3-conda` [Run both C++ and Python tests]
 
 `linux-xenial-gcc5-cuda9-cudnn7-py3-conda` [Run only Python tests for now]
 
@@ -56,12 +56,12 @@ cd TensorComprehensions/docker
 ```
 
 Now, we will build a docker image for TC. We will build docker image for TC on
-linux trusty with gcc4.8, cuda8, cudnn6 and conda environment settings. If you
+linux trusty with gcc4.9, cuda8, cudnn7 and conda environment settings. If you
 need to build images for some other combination, the steps are still valid, just
 make sure to change the paths below to fit your case.
 
 ```Shell
-image=linux-trusty-gcc4.8-cuda8-cudnn7-py3-conda ./docker_build.sh
+image=linux-trusty-gcc4.9-cuda8-cudnn7-py3-conda ./docker_build.sh
 ```
 
 This will build the image for the above permutation and then we can test this image
@@ -74,9 +74,12 @@ run `docker images` and get the `IMAGE_ID` for the image we just built.
 docker run --runtime=nvidia -i -t ${IMAGE_ID}
 # now clone the TC repo
 cd $HOME && git clone https://github.com/facebookresearch/TensorComprehensions.git --recursive && cd TensorComprehensions
+git submodule update --init --recursive
 # build TC
-conda install -y -c pytorch pytorch
-WITH_PYTHON_C2=OFF CLANG_PREFIX=/usr/local/clang+llvm-tapir5.0 ./build.sh --all
+conda install -y mkl-include pyyaml
+conda install -yc conda-forge pytest
+conda install -y pytorch-nightly=2018.04.17 -c pytorch # OR  conda install -y pytorch-nightly=2018.04.17 cuda90 -c pytorch
+CORES=$(nproc) WITH_CAFFE2=ON CLANG_PREFIX=/usr/local/clang+llvm-tapir5.0 BUILD_TYPE=Release ./build.sh --all
 # Test the TC build is fine
 ./test.sh
 ./test_python/run_test.sh

docker/common/install_conda.sh

@@ -18,7 +18,8 @@ popd
 
 export PATH=/opt/conda/bin:$PATH
 
-/opt/conda/bin/conda install numpy decorator six future cmake pyyaml
+/opt/conda/bin/conda install numpy decorator six future cmake pyyaml mkl-include typing
+/opt/conda/bin/conda install -yc conda-forge pytest
 
 which conda
 conda --version

docker/common/install_pip_python.sh

@@ -4,10 +4,17 @@ set -ex
 
 # Install common dependencies
 apt-get update
-apt-get install -y --no-install-recommends python3-dev python3-pip python3-setuptools
+apt-get install -y --no-install-recommends python3-dev python3-setuptools
 
-pip3 install --upgrade pip
-pip3 install numpy decorator six future setuptools pyyaml
+# Install pip from source. The python-pip package on Ubuntu Trusty is old
+# and upon install numpy doesn't use the binary distribution, and fails to compile it from source.
+curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
+python3 get-pip.py
+rm get-pip.py
+
+pip install numpy decorator six future pyyaml
+# install the devel needed for aten
+pip install mkl mkl-devel typing
 
 which python3
 python3 --version

docker/ubuntu-cuda-conda/Dockerfile

@@ -10,8 +10,6 @@ ENV DEBIAN_FRONTEND noninteractive
 # PyTorch conda install
 ARG CUDA_VERSION
 ARG CUDNN_VERSION
-ADD ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh && rm install_cuda.sh
 ENV LD_LIBRARY_PATH /usr/local/cuda/lib64:/usr/local/cuda/targets/x86_64-linux/lib/stubs/:$LD_LIBRARY_PATH
 ENV PATH /usr/local/cuda/bin:/usr/local/bin:$PATH