Skip to content

Simplify copy kernel #28352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zasdfgbnm merged 1 commit into from
Oct 22, 2019
Merged

Simplify copy kernel #28352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions aten/src/ATen/Dispatch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,11 +221,11 @@ inline void deprecated_AT_DISPATCH_ALL_TYPES_AND_HALF_AND_COMPLEX() {}
const auto& SCALAR_TYPE C10_UNUSED = TYPE; \
switch (TYPE) { \
AT_QINT_PRIVATE_CASE_TYPE( \
at::kQInt8, at::qint8, at::kChar, int8_t, __VA_ARGS__) \
at::kQInt8, at::qint8, at::kChar, int8_t, __VA_ARGS__) \
AT_QINT_PRIVATE_CASE_TYPE( \
at::kQUInt8, at::quint8, at::kByte, uint8_t, __VA_ARGS__) \
at::kQUInt8, at::quint8, at::kByte, uint8_t, __VA_ARGS__) \
AT_QINT_PRIVATE_CASE_TYPE( \
at::kQInt32, at::qint32, at::kInt, int, __VA_ARGS__) \
at::kQInt32, at::qint32, at::kInt, int, __VA_ARGS__) \
default: \
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
} \
Expand Down Expand Up @@ -351,6 +351,29 @@ inline void deprecated_AT_DISPATCH_ALL_TYPES_AND_HALF_AND_COMPLEX() {}
} \
}()

#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND_QINTS_AND3(SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, TYPE, NAME, ...) \
[&] { \
switch (TYPE) { \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::ComplexFloat, std::complex<float>, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(at::ScalarType::ComplexDouble, std::complex<double>, __VA_ARGS__) \
AT_QINT_PRIVATE_CASE_TYPE(at::kQInt8, at::qint8, at::kChar, int8_t, __VA_ARGS__) \
AT_QINT_PRIVATE_CASE_TYPE(at::kQUInt8, at::quint8, at::kByte, uint8_t, __VA_ARGS__) \
AT_QINT_PRIVATE_CASE_TYPE(at::kQInt32, at::qint32, at::kInt, int, __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(SCALARTYPE1, decltype(c10::impl::ScalarTypeToCPPType<SCALARTYPE1>::t), __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(SCALARTYPE2, decltype(c10::impl::ScalarTypeToCPPType<SCALARTYPE2>::t), __VA_ARGS__) \
AT_PRIVATE_CASE_TYPE(SCALARTYPE3, decltype(c10::impl::ScalarTypeToCPPType<SCALARTYPE3>::t), __VA_ARGS__) \
default: \
AT_ERROR(#NAME, " not implemented for '", TYPE, "'"); \
} \
}()

// ----------------------------------------------------------------------------
// DEPRECATED MACROS, DON'T USE THESE
// ----------------------------------------------------------------------------
Expand Down
3 changes: 2 additions & 1 deletion aten/src/ATen/native/Copy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,11 +124,12 @@ static Tensor & copy_impl(Tensor & self, const Tensor & src, bool non_blocking)

auto iter = TensorIterator();
iter.set_check_mem_overlap(true);
iter.dont_compute_common_dtype();
iter.add_output(self);
iter.add_input(src);
iter.dont_resize_outputs();
iter.dont_compute_common_dtype();
iter.build();
iter.set_common_dtype(iter.dtype());

if (iter.numel() == 0) {
return self;
Expand Down
4 changes: 4 additions & 0 deletions aten/src/ATen/native/TensorIterator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,10 @@ struct CAFFE2_API TensorIterator {
resize_outputs_ = false;
}

void set_common_dtype(ScalarType dtype) {
common_dtype_ = dtype;
}

void build();

protected:
Expand Down
63 changes: 7 additions & 56 deletions aten/src/ATen/native/cpu/CopyKernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,69 +4,20 @@
#include <ATen/native/Copy.h>
#include <ATen/native/TensorIterator.h>
#include <ATen/native/cpu/Loops.h>
#include <c10/util/TypeCast.h>

namespace at {
namespace native {
namespace {

template <typename self_T>
void copy_kernel_cast(TensorIterator& iter) {
if (isComplexType(iter.dtype(1))) {
AT_DISPATCH_COMPLEX_TYPES(iter.dtype(1), "copy_kernel_cast", [&] {
cpu_kernel(iter, [=](scalar_t a) -> self_T {
return c10::static_cast_with_inter_type<self_T>(std::real(a));
});
});
}
else {
AT_DISPATCH_ALL_TYPES_AND3(
ScalarType::Half,
ScalarType::Bool,
ScalarType::BFloat16,
iter.dtype(1),
"copy_kernel_cast",
[&] {
cpu_kernel(iter, [=](scalar_t a) -> self_T {
return c10::static_cast_with_inter_type<self_T>(a);
});
});
}
}

static void copy_kernel(TensorIterator& iter, bool non_blocking) {
ScalarType dtype = iter.dtype(0);
if (dtype == iter.dtype(1)) {
if (dtype == ScalarType::Half) {
cpu_kernel(iter, [=](at::Half a) -> at::Half { return a; });
} else if (dtype == ScalarType::BFloat16) {
cpu_kernel(iter, [=](at::BFloat16 a) -> at::BFloat16 { return a; });
} else if (isQIntType(dtype)) {
AT_DISPATCH_QINT_TYPES(dtype, "copy_kernel", [&] {
cpu_kernel(
iter,
[=](scalar_t a) -> scalar_t {return a; });
});
} else if (isComplexType(dtype)) {
AT_DISPATCH_COMPLEX_TYPES(dtype, "copy_kernel", [&] {
cpu_kernel(
iter,
[=](scalar_t a) -> scalar_t { return a; });
});
} else {
AT_DISPATCH_ALL_TYPES_AND(
ScalarType::Bool, dtype, "copy_kernel", [&] {
cpu_kernel_vec(
iter,
[=](scalar_t a) -> scalar_t { return a; },
[=](Vec256<scalar_t> a) { return a; });
});
}
} else {
AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(ScalarType::Half, ScalarType::Bool, ScalarType::BFloat16, dtype, "copy_", [&] {
copy_kernel_cast<scalar_t>(iter);
AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND_QINTS_AND3(
ScalarType::Half, ScalarType::Bool, ScalarType::BFloat16, iter.common_dtype(), "copy_",
[&] {
cpu_kernel_vec(
iter,
[](scalar_t a) -> scalar_t { return a; },
[](Vec256<scalar_t> a) { return a; });
});
}
}

} // anonymous namespace
Expand Down
13 changes: 1 addition & 12 deletions aten/src/ATen/native/cuda/Copy.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,12 @@
#include <ATen/native/TensorIterator.h>
#include <ATen/native/cuda/Loops.cuh>
#include <THC/THC.h>
#include <c10/util/TypeCast.h>

namespace at {
namespace native {

using namespace at::cuda;

template <typename dst_t, typename src_t>
void copy_kernel_impl(TensorIterator& iter) {
gpu_kernel(iter, []GPU_LAMBDA(src_t x) -> dst_t {
return c10::static_cast_with_inter_type<dst_t>(x);
});
}

// device-to-device copy, does type conversion
static void copy_device_to_device(TensorIterator& iter, bool non_blocking) {
int64_t numel = iter.numel();
Expand Down Expand Up @@ -67,10 +59,7 @@ static void copy_device_to_device(TensorIterator& iter, bool non_blocking) {
copy_stream));
} else {
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(0), "copy_", [&] {
using dst_t = scalar_t;
AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(1), "copy_", [&] {
copy_kernel_impl<dst_t, scalar_t>(iter);
});
gpu_kernel(iter, []GPU_LAMBDA(scalar_t x) { return x; });
});
}

Expand Down