frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[vulkan] Re-route arithmetic ops to scalar versions when second arg is zero-dim (#73108)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73108

When arithmetic ops are invoked from torchscript the scalar argument will sometimes be wrapped in a zero-dimensional tensor, which will cause the Vulkan implementation to complain as all input tensors are expected to have the same number of channels. The solution is to have the Tensor implementations of the op check if the second argument is zero-dimensional and re-route it to the Scalar implementation if that's the case.

Test Plan: Imported from OSS

Reviewed By: mikaylagawarecki

Differential Revision: D34354840

Pulled By: SS-JIA

fbshipit-source-id: b24799bb3dd4336791a39bea9382c14243ad58e4
(cherry picked from commit c6dd8eb13b9be3800405c64a3a81e5c68da64355)
This commit is contained in:
Sicheng Stephen Jia
2022-02-18 16:53:45 -08:00
committed by PyTorch MergeBot
parent 564f99226a
commit bdc8b3f3e8
1 changed files with 28 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
aten/src/ATen/native/vulkan/ops/Arithmetic.cpp
View File

@ -322,6 +322,13 @@ Tensor add_tensor(
const Tensor& self_arg,
const Tensor& other_arg,
const Scalar& alpha) {
if (other_arg.sizes().size() == 0) {
return arithmetic_scalar(
self_arg,
other_arg.item<float>(),
c10::optional<Scalar>(alpha.to<float>()),
VK_KERNEL(add_scalar));
}
return arithmetic_tensor(
self_arg, other_arg, c10::optional<Scalar>(alpha), VK_KERNEL(add));
}
@ -354,6 +361,13 @@ Tensor sub_tensor(
const Tensor& self_arg,
const Tensor& other_arg,
const Scalar& alpha) {
if (other_arg.sizes().size() == 0) {
return arithmetic_scalar(
self_arg,
other_arg.item<float>(),
c10::optional<Scalar>(-1 * alpha.to<float>()),
VK_KERNEL(add_scalar));
}
return arithmetic_tensor(
self_arg, other_arg, c10::optional<Scalar>(alpha), VK_KERNEL(sub));
}
@ -374,6 +388,13 @@ Tensor& mul_scalar_(Tensor& self, const Scalar& other) {
}
Tensor mul_tensor(const Tensor& self_arg, const Tensor& other_arg) {
if (other_arg.sizes().size() == 0) {
return arithmetic_scalar(
self_arg,
other_arg.item<float>(),
c10::optional<Scalar>(),
VK_KERNEL(mul_scalar));
}
return arithmetic_tensor(
self_arg, other_arg, c10::optional<Scalar>(), VK_KERNEL(mul));
}
@ -400,6 +421,13 @@ Tensor& div_scalar_(Tensor& self, const Scalar& other) {
}
Tensor div_tensor(const Tensor& self_arg, const Tensor& other_arg) {
if (other_arg.sizes().size() == 0) {
return arithmetic_scalar(
self_arg,
1.0 / other_arg.item<float>(),
c10::optional<Scalar>(),
VK_KERNEL(mul_scalar));
}
return arithmetic_tensor(
self_arg, other_arg, c10::optional<Scalar>(), VK_KERNEL(div));
}