Add vulkan shape op (hack)

ghstack-source-id: dc03d44a6151091a0911fe21a493b29b92d287cc Pull Request resolved: #47252
pytorch · Nov 4, 2020 · b2c2b7a · b2c2b7a
1 parent 01da0fe
commit b2c2b7a
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/aten/src/ATen/native/vulkan/ops/Shape.cpp b/aten/src/ATen/native/vulkan/ops/Shape.cpp
@@ -0,0 +1,29 @@
+#include <ATen/native/vulkan/ops/Common.h>
+#include <torch/library.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+namespace ops {
+namespace {
+
+Tensor view(
+    const Tensor& self_arg,
+    IntArrayRef shape) {
+
+  return self_arg.cpu().view(shape).vulkan();
+}
+
+#ifdef USE_VULKAN_API
+
+TORCH_LIBRARY_IMPL(aten, Vulkan, m) {
+  m.impl("view", TORCH_FN(view));
+}
+
+#endif /* USE_VULKAN_API */
+
+} // namespace
+} // namespace ops
+} // namespace vulkan
+} // namespace native
+} // namespace at
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -269,6 +269,26 @@ TEST(VulkanAPITest, avg_pool2d) {
   ASSERT_TRUE(check);
 }
 
+TEST(VulkanAPITest, reshape) {
+  const auto a_cpu = at::rand({1,3,2,2}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto a_vulkan = a_cpu.vulkan();
+
+  const auto c_cpu = at::reshape(a_cpu, {2,3,1,2});
+  const auto c_vulkan = at::reshape(a_vulkan, {2,3,1,2});
+
+  ASSERT_TRUE(almostEqual(c_cpu, c_vulkan.cpu()));
+}
+
+TEST(VulkanAPITest, reshape_) {
+  const auto a_cpu = at::rand({1,3,2,2}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto a_vulkan = a_cpu.vulkan();
+
+  a_cpu.reshape({2,3,1,2});
+  a_vulkan.reshape({2,3,1,2});
+
+  ASSERT_TRUE(almostEqual(a_cpu, a_vulkan.cpu()));
+}
+
 TEST(VulkanAPITest, copy) {
   const auto cpu =
       at::rand({13, 17, 37, 19}, at::device(at::kCPU).dtype(at::kFloat));