Update on "Add acc_gpu_kernel_with_scalars and port add to use it"

ezyang · ezyang · commit 167c9ab95a06 · 2021-08-26T07:41:40.000-07:00
See https://dev-discuss.pytorch.org/t/cuda-loops-case-study-code-generation-vs-templates/302 for explanation of what's going on here. Benchmark script: ``` import torch import torch.utils.benchmark as benchmark results = [] for dtype in [torch.half, torch.float]: for size in [1000, 1000000, 1000000000]: print(f'{dtype} {size}') results.append( benchmark.Timer( stmt='x.add(y)', globals={ 'x': torch.randn(size, dtype=dtype, device='cuda'), 'y': torch.randn(size, dtype=dtype, device='cuda') }, label='perf', sub_label=str(dtype), description=str(size), ).blocked_autorange(min_run_time=1) ) compare = benchmark.Compare(results) compare.print() ``` Quadro GP100 Using the old gpu_kernel_with_scalars: ``` | 1000 | 1000000 | 1000000000 1 threads: -------------------------------------------- torch.float16 | 8.7 | 13.3 | 10796.9 torch.float32 | 8.5 | 24.1 | 21651.8 ``` Quadro GP100 Using the new gpu_kernel_with_scalars (no acc): ``` | 1000 | 1000000 | 1000000000 1 threads: -------------------------------------------- torch.float16 | 8.9 | 13.5 | 10794.7 torch.float32 | 8.8 | 24.0 | 21650.5 Times are in microseconds (us). ``` Quadro GP100 Using the new acc_gpu_kernel_with_scalars: ``` | 1000 | 1000000 | 1000000000 1 threads: -------------------------------------------- torch.float16 | 7.7 | 13.4 | 10795.7 torch.float32 | 7.8 | 23.9 | 21651.6 Times are in microseconds (us). ``` V100 old code: ``` | 100 | 100000 | 100000000 1 threads: ------------------------------------------ torch.int8 | 10.1 | 9.5 | 554.2 torch.int32 | 9.4 | 9.3 | 1411.2 torch.int64 | 9.4 | 9.7 | 3521.1 torch.float16 | 9.5 | 9.5 | 712.1 torch.float32 | 9.6 | 9.6 | 1410.7 ``` V100 new code: ``` | 100 | 100000 | 100000000 1 threads: ------------------------------------------ torch.int8 | 10.0 | 9.6 | 554.1 torch.int32 | 10.1 | 9.9 | 1411.2 torch.int64 | 10.5 | 10.5 | 3520.9 torch.float16 | 10.4 | 10.2 | 712.1 torch.float32 | 10.1 | 10.0 | 1411.3 ``` Signed-off-by: Edward Z. Yang <ezyang@fb.com> Differential Revision: [D30545296](https://our.internmc.facebook.com/intern/diff/D30545296) [ghstack-poisoned]
diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh
@@ -5,6 +5,7 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/TensorIteratorDynamicCasting.h>
 #include <ATen/cuda/detail/OffsetCalculator.cuh>
+#include <ATen/OpMathType.h>
 
 #include <thrust/tuple.h>
 
