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

Co-authored-by: Xuehai Pan <XuehaiPan@pku.edu.cn>

Author: JieRen98

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).

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;
   }
 }