mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
517c7c9861
Summary:
Anywhere we used #include "foo.h", we now say #include <foo.h>
Paths are adjusted to be rooted out of aten/src, torch/lib, or
the root level directory.
I modified CMakeLists.txt by hand to remove TH and THC from
the include paths.
I used the following script to do the canonicalization:
```
import subprocess
import re
import os.path
files = subprocess.check_output(['git', 'ls-files']).decode('utf-8').rstrip().split('\n')
for fn in files:
if not any(fn.endswith(suff) for suff in ['.cu', '.cpp', '.in', '.h', '.hpp', '.cu', '.cuh', '.cc']):
continue
if not any(fn.startswith(pref) for pref in ["aten/", "torch/"]):
continue
with open(fn, 'r') as f:
c = f.read()
def fmt(p):
return "#include <{}>".format(p)
def repl(m):
p = m.group(1)
if p in ["dlfcn.h", "unistd.h", "nvrtc.h", "cuda.h", "cuda_runtime.h", "cstdint", "cudnn.h", "Python.h", "cusparse.h", "cuda_runtime_api.h", "cuda_fp16.h", "cublas_v2.h", "stdint.h", "curand_kernel.h"]:
return fmt(p)
if any(p.startswith(pref) for pref in ["torch/csrc", "c10/", "ATen/", "caffe2/", "TH/", "THC/", "Eigen/", "gtest/", "zdl/", "gloo/", "onnx/", "miopen/"]):
return fmt(p)
for root in ["aten/src", "torch/lib", ""]:
for bad_root in [os.path.dirname(fn), "aten/src/TH", "aten/src/THC", "torch/csrc"]:
new_p = os.path.relpath(os.path.join(bad_root, p), root)
if not new_p.startswith("../") and (os.path.exists(os.path.join(root, new_p)) or os.path.exists(os.path.join(root, new_p + ".in"))):
return fmt(new_p)
print("ERROR: ", fn, p)
return m.group(0)
new_c = re.sub(r'#include "([^"]+)"', repl, c)
if new_c != c:
print(fn)
with open(fn, 'w') as f:
f.write(new_c)
```
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14849
Reviewed By: dzhulgakov
Differential Revision: D13363445
Pulled By: ezyang
fbshipit-source-id: 52361f878a672785f9306c9e9ab2513128092b68
100 lines
3.1 KiB
C++
100 lines
3.1 KiB
C++
#include <torch/csrc/utils/tensor_apply.h>
|
|
|
|
#include <ATen/TensorUtils.h>
|
|
#include <ATen/ExpandUtils.h>
|
|
|
|
#include <torch/csrc/Exceptions.h>
|
|
#include <torch/csrc/utils/python_numbers.h>
|
|
#include <torch/csrc/utils/python_scalars.h>
|
|
|
|
using namespace at;
|
|
|
|
namespace torch { namespace utils {
|
|
|
|
struct StridedData {
|
|
StridedData(const Tensor & tensor)
|
|
: data(tensor.data_ptr())
|
|
, strides(tensor.strides())
|
|
, elementSize(tensor.type().elementSizeInBytes()) {}
|
|
|
|
void* data;
|
|
IntList strides;
|
|
int64_t elementSize;
|
|
|
|
void step(int dim) {
|
|
data = (char*)data + (strides[dim] * elementSize);
|
|
}
|
|
};
|
|
|
|
template<size_t N>
|
|
static void recursive_apply(IntList sizes, ScalarType scalarType, int64_t dim,
|
|
PyObject* fn, std::array<StridedData, N> strided_data) {
|
|
int64_t ndim = sizes.size();
|
|
if (dim == ndim) {
|
|
auto args = THPObjectPtr(PyTuple_New(N));
|
|
if (!args) throw python_error();
|
|
for (size_t i = 0; i < N; i++) {
|
|
PyObject* arg = load_scalar(strided_data[i].data, scalarType);
|
|
if (!arg) throw python_error();
|
|
PyTuple_SET_ITEM(args.get(), i, arg);
|
|
}
|
|
auto ret = THPObjectPtr(PyObject_CallObject(fn, args.get()));
|
|
if (!ret) throw python_error();
|
|
store_scalar(strided_data[0].data, scalarType, ret.get());
|
|
return;
|
|
}
|
|
|
|
auto n = sizes[dim];
|
|
for (int64_t i = 0; i < n; i++) {
|
|
recursive_apply(sizes, scalarType, dim + 1, fn, strided_data);
|
|
for (auto& td : strided_data) {
|
|
td.step(dim);
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor & apply_(Tensor & self, PyObject* fn) {
|
|
if (self.type().backend() != Backend::CPU) {
|
|
throw TypeError("apply_ is only implemented on CPU tensors");
|
|
}
|
|
auto scalarType = self.type().scalarType();
|
|
recursive_apply<1>(self.sizes(), scalarType, 0, fn, {{ self }});
|
|
return self;
|
|
}
|
|
|
|
Tensor & map_(Tensor & self, const Tensor & other_, PyObject* fn) {
|
|
if (self.type().backend() != Backend::CPU) {
|
|
throw TypeError("map_ is only implemented on CPU tensors");
|
|
}
|
|
if (other_.type() != self.type()) {
|
|
throw TypeError("map_: expected %s for 'other' (got %s)",
|
|
self.type().toString(), other_.type().toString());
|
|
}
|
|
Tensor other;
|
|
std::tie(other) = expand_inplace(self, other_, "map_");
|
|
auto scalarType = self.type().scalarType();
|
|
recursive_apply<2>(self.sizes(), scalarType, 0, fn, {{ self, other }});
|
|
return self;
|
|
}
|
|
|
|
Tensor & map2_(Tensor & self, const Tensor & x_, const Tensor & y_, PyObject* fn) {
|
|
if (self.type().backend() != Backend::CPU || x_.type().backend() != Backend::CPU || y_.type().backend() != Backend::CPU) {
|
|
throw TypeError("map2_ is only implemented on CPU tensors");
|
|
}
|
|
if (x_.type() != self.type()) {
|
|
throw TypeError("map2_: expected %s for argument 'x' (got %s)",
|
|
self.type().toString(), x_.type().toString());
|
|
}
|
|
if (y_.type() != self.type()) {
|
|
throw TypeError("map2_: expected %s for argument 'y' (got %s)",
|
|
self.type().toString(), y_.type().toString());
|
|
}
|
|
Tensor other1, other2;
|
|
std::tie(other1, other2) = expand_inplace(self, x_, y_, "map2_");
|
|
auto scalarType = self.type().scalarType();
|
|
recursive_apply<3>(self.sizes(), scalarType, 0, fn, {{ self, other1, other2 }});
|
|
return self;
|
|
}
|
|
|
|
}} // namespace torch::utils
|