[Performance] Use allocator from PyTorch if possible (#2328)

* first commit

* some thoughts

* move around

* more commit

* more fixes

* now it uses torch allocator

* fix symbol export error

* fix

* fixes

* test fix

* add script

* building separate library per version

* fix for vs2019

* more fixes

* fix on windows build

* update jenkinsfile

* auto copy built dlls for windows

* lint and installation guide update

* fix

* specify conda environment

* set environment for ci

* fix

* fix

* fix

* fix again

* revert

* fix cmake

* fix

* switch to using python interpreter path

* remove scripts

* debug

* oops sorry

* Update index.rst

* Update index.rst

* copies automatically, no need for this

* do not print message if library not found

* tiny fixes

* debug on nightly

* replace add_compile_definitions to make CMake 3.5 happy

* fix linking to wrong lib for multiple pytorch envs

* changed building strategy

* fix nightly

* fix windows

* fix windows again

* setup bugfix

* address comments

* change README

tensoradapter/pytorch/build.bat

+REM Helper script to build tensor adapter libraries for PyTorch
+@ECHO OFF
+SETLOCAL EnableDelayedExpansion
+
+MD "%BINDIR%\tensoradapter\pytorch"
+DEL /S /Q build
+MD build
+PUSHD build
+
+IF x%1x == xx GOTO single
+
+FOR %%X IN (%*) DO (
+	DEL /S /Q *
+	"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
+	msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
+	COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
+)
+
+GOTO end
+
+:single
+
+DEL /S /Q *
+"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" .. -G "Visual Studio 16 2019" || EXIT /B 1
+msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
+COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
+
+:end
+POPD
+
+ENDLOCAL

tensoradapter/pytorch/build.sh

+#!/bin/bash
+# Helper script to build tensor adapter libraries for PyTorch
+set -e
+
+rm -rf build
+mkdir -p build
+mkdir -p $BINDIR/tensoradapter/pytorch
+cd build
+
+if [ $# -eq 0 ]; then
+	${CMAKE_COMMAND} -DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR ..
+	make -j
+	cp -v *.so $BINDIR/tensoradapter/pytorch
+else
+	for PYTHON_INTERP in $@; do
+		rm -rf *
+		${CMAKE_COMMAND} -DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DPYTHON_INTERP=$PYTHON_INTERP ..
+		make -j
+		cp -v *.so $BINDIR/tensoradapter/pytorch
+	done
+fi

tensoradapter/pytorch/find_cmake.py

+import torch
+import os
+
+cmake_prefix_path = getattr(
+    torch.utils,
+    "cmake_prefix_path",
+    os.path.join(os.path.dirname(torch.__file__), "share", "cmake"))
+version = torch.__version__.split('+')[0]
+print(';'.join([cmake_prefix_path, version]))

tensoradapter/pytorch/torch.cpp

+/*!
+ *  Copyright (c) 2020 by Contributors
+ * \file torch/torch.cpp
+ * \brief Implementation of PyTorch adapter library.
+ */
+
+#include <tensoradapter.h>
+#include <torch/torch.h>
+#include <ATen/DLConvertor.h>
+#include <vector>
+#include <iostream>
+
+namespace tensoradapter {
+
+static at::Device get_device(DLContext ctx) {
+  switch (ctx.device_type) {
+   case kDLCPU:
+    return at::Device(torch::kCPU);
+    break;
+   case kDLGPU:
+    return at::Device(torch::kCUDA, ctx.device_id);
+    break;
+   default:
+    // fallback to CPU
+    return at::Device(torch::kCPU);
+    break;
+  }
+}
+
+extern "C" {
+
+DLManagedTensor* TAempty(
+    std::vector<int64_t> shape,
+    DLDataType dtype,
+    DLContext ctx) {
+  auto options = torch::TensorOptions()
+    .layout(torch::kStrided)
+    .device(get_device(ctx))
+    .dtype(at::toScalarType(dtype));
+  torch::Tensor tensor = torch::empty(shape, options);
+  return at::toDLPack(tensor);
+}
+
+};
+
+};  // namespace tensoradapter

tests/scripts/build_dgl.bat

 @ECHO OFF
 SETLOCAL EnableDelayedExpansion
 
+CALL "C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
 CALL mkvirtualenv --system-site-packages %BUILD_TAG%
 DEL /S /Q build
 DEL /S /Q _download
 MD build
 
 SET _MSPDBSRV_ENDPOINT_=%BUILD_TAG%
-SET TMP=%WORKSPACE%\\tmp
-SET TEMP=%WORKSPACE%\\tmp
-SET TMPDIR=%WORKSPACE%\\tmp
+SET TMP=%WORKSPACE%\tmp
+SET TEMP=%WORKSPACE%\tmp
+SET TMPDIR=%WORKSPACE%\tmp
 
 PUSHD build
-CALL "C:\Program Files (x86)\Microsoft Visual Studio\2017\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
-cmake -DCMAKE_CXX_FLAGS="/DDGL_EXPORTS" -DUSE_OPENMP=ON -Dgtest_force_shared_crt=ON -DDMLC_FORCE_SHARED_CRT=ON -DBUILD_CPP_TEST=1 -DCMAKE_CONFIGURATION_TYPES="Release" .. -G "Visual Studio 15 2017 Win64" || EXIT /B 1
+cmake -DCMAKE_CXX_FLAGS="/DDGL_EXPORTS" -DUSE_OPENMP=ON -DBUILD_TORCH=ON -Dgtest_force_shared_crt=ON -DDMLC_FORCE_SHARED_CRT=ON -DBUILD_CPP_TEST=1 -DCMAKE_CONFIGURATION_TYPES="Release" .. -G "Visual Studio 16 2019" || EXIT /B 1
 msbuild dgl.sln /m /nr:false || EXIT /B 1
-COPY Release\dgl.dll .
-COPY Release\runUnitTests.exe .
+COPY /Y Release\runUnitTests.exe .
 POPD
 
+CALL workon %BUILD_TAG%
+
 PUSHD python
 DEL /S /Q build *.egg-info dist
 pip install -e . || EXIT /B 1

tests/scripts/build_dgl.sh

@@ -7,7 +7,11 @@ if [ $# -ne 1 ]; then
     exit -1
 fi
 
-CMAKE_VARS="-DBUILD_CPP_TEST=ON -DUSE_OPENMP=ON"
+CMAKE_VARS="-DBUILD_CPP_TEST=ON -DUSE_OPENMP=ON -DBUILD_TORCH=ON"
+# This is a semicolon-separated list of Python interpreters containing PyTorch.
+# The value here is for CI.  Replace it with your own or comment this whole
+# statement for default Python interpreter.
+CMAKE_VARS="$CMAKE_VARS -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
 
 if [ "$1" == "gpu" ]; then
     CMAKE_VARS="-DUSE_CUDA=ON $CMAKE_VARS"
@@ -36,4 +40,4 @@ python3 setup.py install
 # test inplace build (for cython)
 python3 setup.py build_ext --inplace
 done
-popd
+popd
\ No newline at end of file