Link with pytorch/xla when build for Linux

.circleci/config.yml

@@ -9,7 +9,8 @@ jobs:
     steps:
       - checkout
       - run: pre-commit run -a
-      - run: LD_LIBRARY_PATH=$CIRCLE_WORKING_DIRECTORY/cgotorch/libtorch/lib
+      # yamllint disable-line rule:line-length
+      - run: export LD_LIBRARY_PATH=$CIRCLE_WORKING_DIRECTORY/cgotorch/libtorch/lib:$CIRCLE_WORKING_DIRECTORY/cgotorch/torch_xla/lib:$LD_LIBRARY_PATH
       - run: go generate ./...
       - run: go install ./...
       - run: go test -race -coverprofile=coverage.txt -covermode=atomic -v ./...

.gitignore

@@ -1,4 +1,3 @@
-libtorch*
 build
 gotorch
 libcgotorch.so

benchmark/Makefile

@@ -8,6 +8,7 @@ mnist: main.cpp
 	main.cpp \
 	-O -o mnist \
 	-Wl,-rpath,${GOPATH}/src/github.com/wangkuiyi/gotorch/cgotorch/libtorch/lib \
+	-Wl,-rpath,${GOPATH}/src/github.com/wangkuiyi/gotorch/cgotorch/torch_xla/lib \
 	-pthread \
 	-lc10 -ltorch -ltorch_cpu \
 	-D_GLIBCXX_USE_CXX11_ABI=0

cgotorch/.gitignore

@@ -0,0 +1,4 @@
+libtorch*
+torch_xla
+*.bz2
+*.zip

cgotorch/Makefile

@@ -16,8 +16,9 @@ libcgotorch.so : $(srcs) $(hdrs) ${LIBTORCH_DIR}
 	$(srcs) \
 	-O -o $@ ${INSTALL_NAME} \
 	-Wl,-rpath,libtorch/lib \
+	-Wl,-rpath,torch_xla/lib \
 	-Wl,-${LOAD} libtorch/lib/libc10.${LIB_SUFFIX} \
-	-lc10 -ltorch -ltorch_cpu \
+	-lc10 -ltorch -ltorch_cpu ${XLA_LIBS} \
 	-D_GLIBCXX_USE_CXX11_ABI=${GLIBCXX_USE_CXX11_ABI}
 
 clean:

cgotorch/build.sh

@@ -2,70 +2,106 @@
 
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 
-pushd $DIR
+pushd $DIR > /dev/null
 
-OS=$(uname | tr '[:upper:]' '[:lower:]')
-ARCH=$(uname -m)
 CXX="g++"
 LIBTORCH_DIR=""
 GLIBCXX_USE_CXX11_ABI="1"
 LOAD="force_load"
 LIB_SUFFIX="so"
 INSTALL_NAME=""
+XLA_LIBS=""
 
-if [[ "$OS" == "linux" ]]; then
-    if [[ "$ARCH" =~ arm* ]]; then
-        echo "Building for Raspbian ...";
-        LIBTORCH_DIR="rpi/libtorch"
-        if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
-            curl -LsO 'https://github.com/ljk53/pytorch-rpi/raw/master/libtorch-rpi-cxx11-abi-shared-1.6.0.zip';
-            unzip -qq -o libtorch-rpi-cxx11-abi-shared-1.6.0.zip -d rpi
-        fi
-    elif $(whereis cuda | cut -f 2 -d ' ')/bin/nvcc --version > /dev/null; then
-        CXX="clang++"
-        CUDA_VERSION=`nvcc --version | grep release | grep -Eo "[0-9]+.[0-9]+" | head -1`
-        if [[ "$CUDA_VERSION" == "10.1" ]]; then
-            echo "Building for Linux with CUDA 10.1";
-            LIBTORCH_DIR="linux-cuda101/libtorch"
-            if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
-                curl -Lso libtorch-cxx11-1.6.0-linux-cuda101.zip 'https://download.pytorch.org/libtorch/cu101/libtorch-cxx11-abi-shared-with-deps-1.6.0%2Bcu101.zip'
-                unzip -qq -o libtorch-cxx11-1.6.0-linux-cuda101.zip -d linux-cuda101
-            fi
-        elif [[ "$CUDA_VERSION" == "10.2" ]]; then
-            echo "Building for Linux with CUDA 10.2";
-            LIBTORCH_DIR="linux-cuda102/libtorch"
-            if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
-                curl -Lso libtorch-cxx11-1.6.0-linux-cuda102.zip 'https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.6.0.zip'
-                unzip -qq -o libtorch-cxx11-1.6.0-linux-cuda102.zip -d linux-cuda102
-            fi
-        fi
-    else
-        echo "Building for Linux without CUDA ...";
-        LIBTORCH_DIR="linux/libtorch"
-        GLIBCXX_USE_CXX11_ABI="0"
-        if [[ ! -d "DIR/$LIBTORCH_DIR" ]]; then
-            curl -LsO 'https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-1.6.0%2Bcpu.zip'
-            unzip -qq -o libtorch-shared-with-deps-1.6.0%2Bcpu.zip -d linux
-        fi
-    fi
-elif [[ "$OS" == "darwin" ]]; then
+function build_macos() {
     echo "Building for macOS ...";
     LIBTORCH_DIR="macos/libtorch"
     LIB_SUFFIX="dylib"
     INSTALL_NAME="-install_name @rpath/\$@"
     LOAD="all_load"
     if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
-        curl -LsO https://download.pytorch.org/libtorch/cpu/libtorch-macos-1.6.0.zip
+        curl -LsO "https://download.pytorch.org/libtorch/cpu/libtorch-macos-1.6.0.zip"
         unzip -qq -o libtorch-macos-1.6.0.zip -d macos
     fi
+}
+
+function build_raspbian() {
+    echo "Building for Raspbian ...";
+    LIBTORCH_DIR="rpi/libtorch"
+    if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
+        curl -LsO "https://github.com/ljk53/pytorch-rpi/raw/master/libtorch-rpi-cxx11-abi-shared-1.6.0.zip"
+        unzip -qq -o libtorch-rpi-cxx11-abi-shared-1.6.0.zip -d rpi
+    fi
+}
+
+function build_cuda() {
+    CUDA_VERSION=$("$NVCC" --version | grep release | grep -Eo "[0-9]+.[0-9]+" | head -1)
+
+    if [[ "$CUDA_VERSION" == "10.1" ]]; then
+        CUDA_VERSION="101"
+    elif [[ "$CUDA_VERSION" == "10.2" ]]; then
+        CUDA_VERSION="102"
+    else
+        echo "Not supported CUDA version. $CUDA_VERSION"
+        return -1
+    fi
+
+    echo "Building for Linux with CUDA 10.1";
+    CXX="clang++"
+
+    LIBTORCH_DIR="linux-cuda$CUDA_VERSION/libtorch"
+    if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
+        curl -Lso libtorch-cxx11-1.6.0-linux-cuda$CUDA_VERSION.zip 'https://download.pytorch.org/libtorch/cu$CUDA_VERSION/libtorch-cxx11-abi-shared-with-deps-1.6.0%2Bcu$CUDA_VERSION.zip'
+        unzip -qq -o libtorch-cxx11-1.6.0-linux-cuda$CUDA_VERSION.zip -d linux-cuda$CUDA_VERSION
+    fi
+}
+
+function build_linux(){
+    echo "Building for Linux without CUDA ...";
+    LIBTORCH_DIR="linux/libtorch"
+    GLIBCXX_USE_CXX11_ABI="0"
+    if [[ ! -d "$DIR/$LIBTORCH_DIR" ]]; then
+        curl -LsO 'https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-1.6.0%2Bcpu.zip'
+        unzip -qq -o libtorch-shared-with-deps-1.6.0%2Bcpu.zip -d linux
+    fi
+
+    if [[ ! -d "$DIR/torch_xla" ]]; then
+        curl -LsO "https://github.com/wangkuiyi/torch_xla_prebuilt/raw/master/torch_xla-1.6.0.tar.bz2"
+        tar xjf torch_xla-1.6.0.tar.bz2
+    fi
+
+    if [[ -d "$DIR/torch_xla" ]]; then
+        echo "Found XLA library. Use it."
+        XLA_LIBS="-Ltorch_xla/lib -lptxla -lxla_computation_client"
+    fi
+}
+
+
+OS=$(uname | tr '[:upper:]' '[:lower:]')
+ARCH=$(uname -m)
+NVCC=$(whereis cuda | cut -f 2 -d ' ')"/bin/nvcc"
+
+if [[ "$OS" == "darwin" ]]; then
+    build_macos
+elif [[ "$OS" == "linux" ]]; then
+    if [[ "$ARCH" =~ arm* ]]; then
+        build_raspbian
+    elif "$NVCC" --version > /dev/null; then
+        if ! build_cuda; then
+            build_linux
+        fi
+    else
+        build_linux
+    fi
 fi
 
+
 make CXX="$CXX" \
      LIB_SUFFIX="$LIB_SUFFIX" \
      INSTALL_NAME="$INSTALL_NAME" \
      LIBTORCH_DIR="$LIBTORCH_DIR" \
      GLIBCXX_USE_CXX11_ABI="$GLIBCXX_USE_CXX11_ABI" \
      LOAD="$LOAD" \
+     XLA_LIBS="$XLA_LIBS" \
      -f Makefile;
 
-popd
+popd > /dev/null

device.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 import "unsafe"

nn/functional/functional.go

@@ -3,6 +3,7 @@ package functional
 // #cgo CFLAGS: -I ${SRCDIR}/../../cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch -Wl,-rpath ${SRCDIR}/../../cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/../../cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/../../cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "../../cgotorch/cgotorch.h"
 import "C"
 

nn/initializer/initializer.go

@@ -3,6 +3,7 @@ package initializer
 // #cgo CFLAGS: -I ${SRCDIR}/../../cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch -Wl,-rpath ${SRCDIR}/../../cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/../../cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/../../cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/../../cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "../../cgotorch/cgotorch.h"
 import "C"
 

optim.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 import (

tensor.go

@@ -5,6 +5,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 

tensor_ctor.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 

tensor_gc.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 import (

tensor_gob.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"
 

tensor_ops.go

@@ -3,6 +3,7 @@ package gotorch
 // #cgo CFLAGS: -I ${SRCDIR}/cgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch -Wl,-rpath ${SRCDIR}/cgotorch -lcgotorch
 // #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/libtorch/lib -Wl,-rpath ${SRCDIR}/cgotorch/libtorch/lib -lc10 -ltorch -ltorch_cpu
+// #cgo LDFLAGS: -L ${SRCDIR}/cgotorch/torch_xla/lib -Wl,-rpath ${SRCDIR}/cgotorch/torch_xla/lib -lptxla -lxla_computation_client
 // #include "cgotorch.h"
 import "C"