perf(acc_op): further acceleration with CUDA unroll (#112)

JieRen98 · XuehaiPan · web-flow · commit b155b1523e51 · 2022-12-06T20:57:29.000+08:00
Co-authored-by: Xuehai Pan &lt;XuehaiPan@pku.edu.cn&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add unroll pragma for CUDA OPs by [@JieRen98](https://github.com/JieRen98) and [@XuehaiPan](https://github.com/XuehaiPan) in [#112](https://github.com/metaopt/torchopt/pull/112).
 - Add Python implementation of accelerated OP and pure-Python wheels by [@XuehaiPan](https://github.com/XuehaiPan) in [#67](https://github.com/metaopt/torchopt/pull/67).
 - Add `nan_to_num` hook and gradient transformation by [@XuehaiPan](https://github.com/XuehaiPan) in [#119](https://github.com/metaopt/torchopt/pull/119).
 - Add matrix inversion linear solver with neumann series approximation by [@Benjamin-eecs](https://github.com/Benjamin-eecs) and [@XuehaiPan](https://github.com/XuehaiPan) in [#98](https://github.com/metaopt/torchopt/pull/98).
diff --git a/src/adam_op/adam_op_impl_cpu.cpp b/src/adam_op/adam_op_impl_cpu.cpp
@@ -135,7 +135,7 @@ void adamForwardNuCPUKernel(const scalar_t *__restrict__ updates_ptr,
     const scalar_t updates = updates_ptr[tid];
     const scalar_t nu = nu_ptr[tid];
 
-    const scalar_t nu_out = b2 * nu + (1 - b2) * pow(updates, 2);
+    const scalar_t nu_out = b2 * nu + (1 - b2) * updates * updates;
     nu_out_ptr[tid] = nu_out;
   }
 }
diff --git a/src/adam_op/adam_op_impl_cuda.cu b/src/adam_op/adam_op_impl_cuda.cu

Original file line number	Diff line number	Diff line change
`@@ -135,7 +135,7 @@ void adamForwardNuCPUKernel(const scalar_t *__restrict__ updates_ptr,`
`135`	`135`	`const scalar_t updates = updates_ptr[tid];`
`136`	`136`	`const scalar_t nu = nu_ptr[tid];`
`137`	`137`
`138`		`- const scalar_t nu_out = b2 * nu + (1 - b2) * pow(updates, 2);`
	`138`	`+ const scalar_t nu_out = b2 * nu + (1 - b2) * updates * updates;`
`139`	`139`	`nu_out_ptr[tid] = nu_out;`
`140`	`140`	`}`
`141`	`141`	`}`