Add AMP support to linalg.vecdot. (pytorch#108165)

We follow the same rules as matmul. Fixes pytorch#108127 Pull Request resolved: pytorch#108165 Approved by: https://github.com/albanD
sands-lab · Aug 29, 2023 · 86bc50a · 86bc50a
1 parent 75884f4
commit 86bc50a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/aten/src/ATen/autocast_mode.cpp b/aten/src/ATen/autocast_mode.cpp
@@ -273,6 +273,7 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) {
   KERNEL_CUDA(einsum, lower_precision_fp)
   KERNEL_CUDA(mm, lower_precision_fp)
   KERNEL_CUDA(mv, lower_precision_fp)
+  KERNEL_CUDA(linalg_vecdot, lower_precision_fp)
   KERNEL_CUDA(linear, lower_precision_fp)
   KERNEL_CUDA(addbmm, lower_precision_fp)
   KERNEL_CUDA(baddbmm, lower_precision_fp)
@@ -395,6 +396,7 @@ TORCH_LIBRARY_IMPL(aten, AutocastCPU, m) {
   KERNEL_CPU2(conv3d, padding, lower_precision_fp)
   KERNEL_CPU(bmm, lower_precision_fp)
   KERNEL_CPU(mm, lower_precision_fp)
+  KERNEL_CPU(linalg_vecdot, lower_precision_fp)
   KERNEL_CPU(baddbmm, lower_precision_fp)
   KERNEL_CPU(addmm, lower_precision_fp)
   KERNEL_CPU(addbmm, lower_precision_fp)

diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py
@@ -217,6 +217,7 @@ def __init__(self, dev):
             ("multi_margin_loss", mat0_fp16 + (torch.ones((n,), device=dev, dtype=torch.long),)),
         ]
         self.linalg_fp16 = [
+            ("linalg_vecdot", mat0_fp32 + mat0_fp32),
             ("linalg_multi_dot", (mat0_fp32 + mat1_fp32 + mat2_fp32,)),
         ]
         self.methods_fp16 = [