dmlc--dgl
Quan (Andy) Gan
9a7235faf2
* 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
47 行
954 B
C++
47 行
954 B
C++
/*!
|
|
* 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
|