Use gcc 5.4 on xenial non-conda + use pytorch 0.4 official + fix python tests

fix the aten build which fails with avx related issues on gcc5.5. Ubuntu toolchain ships gcc 5.5 by default since April 24, 2018 hence not using the toolchain for gcc 5.4 but rather what ubuntu provides. Also use the official pytorch 0.4 version and fix python tests which were broken by a recent change in PR 353

Author: prigoyal

.jenkins/build.sh

@@ -72,7 +72,7 @@ if [[ "$DISTRIB_RELEASE" == 14.04 ]]; then
     source activate tc-env
     conda install -y pyyaml mkl-include
     conda install -yc conda-forge pytest
-    conda install -y pytorch-nightly=2018.04.17 -c pytorch
+    conda install -y pytorch -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"
@@ -87,7 +87,7 @@ if [[ "$DISTRIB_RELEASE" == 16.04 ]]; then
     source activate tc-env
     conda install -y pyyaml mkl-include
     conda install -yc conda-forge pytest
-    conda install -y pytorch-nightly=2018.04.17 cuda90 -c pytorch
+    conda install -y pytorch 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"

docker/README.md

@@ -78,7 +78,7 @@ git submodule update --init --recursive
 # build TC
 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
+conda install -y pytorch -c pytorch # OR  conda install -y pytorch 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

docker/common/install_base.sh

@@ -34,11 +34,15 @@ apt-get install -y --no-install-recommends \
 apt-get clean
 rm -rf /var/lib/apt/lists/*
 # setup gcc
-add-apt-repository ppa:ubuntu-toolchain-r/test
-apt-get update
-apt-get install -y --no-install-recommends libcilkrts5 gcc-$GCC_VERSION g++-$GCC_VERSION
-update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-$GCC_VERSION 50
-update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-$GCC_VERSION 50
+if [[ "$GCC_VERSION" == 4.9 ]]; then
+  add-apt-repository ppa:ubuntu-toolchain-r/test
+  apt-get update
+  apt-get install -y --no-install-recommends libcilkrts5 gcc-$GCC_VERSION g++-$GCC_VERSION
+  update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-$GCC_VERSION 50
+  update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-$GCC_VERSION 50
+else
+  apt-get install -y --no-install-recommends libcilkrts5 gcc g++
+fi
 
 # Install ccache from source. Needs 3.4 or later for ccbin support
 # Needs specific branch to work with nvcc (ccache/ccache#145)

tensor_comprehensions/tc_unit.py

@@ -66,7 +66,7 @@ def set_gflags(
 def check_cache_file_exists(cache_file):
     # for autotuning, we save two files: .cuda and .options, we will check that
     # these two files exists for the validity of cache
-    if os.path.exists(cache_file + ".options") and os.path.exists(cache_file + ".cuda"):
+    if os.path.exists(cache_file + ".options"):
         return True
     return False
 
@@ -254,7 +254,7 @@ def autotune(self, *inputs, **kwargs):
                 cache_file = "/tmp/{}_{}".format(hash_key, str(uuid.uuid4()))
             elif isinstance(kwargs["cache"], str):
                 cache_file = kwargs["cache"]
-            logger.info('Autotuning cache will be saved to: {}.cuda/options'.format(cache_file))
+            logger.info('Autotuning cache will be saved to: {}.options'.format(cache_file))
         else:
             logger.warning("Autotuning results won't be cached. 'cache' option is not set")
 
@@ -297,7 +297,7 @@ def autotune(self, *inputs, **kwargs):
 
         if cache_file:
             cache_file = cache_file + "_backward"
-            logger.info('Backwards autotuning cache will be saved to: {}.cuda/options'.format(cache_file))
+            logger.info('Backwards autotuning cache will be saved to: {}.options'.format(cache_file))
         kwargs["type"] = "backward"
         options = get_options_from_kwargs_and_tuner_cache(backward_name, cache_file, options_cache, *inputs, **kwargs)
         backward_best_options = self.tune_and_store(

test_python/test_tc.py

@@ -79,7 +79,9 @@ def matmul(float(M,N) A, float(N,K) B) -> (output) {
         inputs = [mat1, mat2]
         handle = cu.compile("matmul", [mat1, mat2], options="mlp")
         outputs = cu.run(handle, "matmul", inputs)
+        torch.cuda.synchronize()
         expected = torch.mm(mat1, mat2)
+        torch.cuda.synchronize()
         diff = outputs[0] - expected
         self.assert_almost_equal(diff, inputs, 4)
 

test_python/test_tc_torch.py

@@ -181,7 +181,7 @@ def test_autotuner_cachefile_first(self):
     def test_autotuner_cachefile_load_automatic(self):
         lang = MATMUL_LANG
         cache_file = "{}/matmul_100_400_500".format(PATH_PREFIX)    # use argparse if input from command line
-        assert os.path.isfile("{}.cuda".format(cache_file)), "looks like the cache_file doesn't exist"
+        assert os.path.isfile("{}.options".format(cache_file)), "looks like the cache_file doesn't exist"
 
         matmul = tc.define(lang, name="matmul")
         mat1, mat2 = torch.randn(100, 400).cuda(), torch.randn(400, 500).cuda()