add fused adam (#822)

* add fused adam

* remove unused code

* remove the fix for launcher script out of this PR

* move all scalar value cast to the top of the function

* using more meaningful variable names

Author: zhuhaozhe

intel_extension_for_pytorch/csrc/aten/cpu/kernels/optimizer/AdamFusedStepKrnl.cpp

@@ -0,0 +1,117 @@
+#include "optimizer.h"
+
+#include <torch/csrc/autograd/function.h>
+#include <torch/extension.h>
+#include "csrc/utils/ipex_op_profile.h"
+
+namespace torch_ipex {
+namespace cpu {
+
+DEFINE_DISPATCH(adam_fused_step_kernel_stub);
+
+void adam_fused_step(
+    const at::Tensor& param_,
+    const at::Tensor& exp_avg_,
+    const at::Tensor& exp_avg_sq_,
+    const at::Tensor& max_exp_avg_sq_,
+    const at::Tensor& grad_,
+    const at::Tensor& param2_,
+    bool amsgrad,
+    double step,
+    double beta1,
+    double beta2,
+    double learning_rate,
+    double weight_decay,
+    double eps) {
+  IPEX_RECORD_FUNCTION(
+      "torch_ipex::adam_fused_step", c10::ArrayRef<c10::IValue>({}));
+
+  TORCH_CHECK(
+      learning_rate >= 0, "Expect learning rate >= 0.0, got ", learning_rate);
+  TORCH_CHECK(eps >= 0, "Expect eps >= 0.0, got ", eps);
+  TORCH_CHECK(beta1 >= 0 && beta1 < 1, "Expect 0.0 <= beta1 < 1.0, got", beta1);
+  TORCH_CHECK(beta2 >= 0 && beta2 < 1, "Expect 0.0 <= beta2 < 1.0, got", beta2);
+  TORCH_CHECK(
+      weight_decay >= 0, "Expect weight_decay >= 0.0, got ", weight_decay);
+
+  TORCH_CHECK(
+      param_.sizes() == grad_.sizes(),
+      "Expect param and grad have the same sizes, param sizes: ",
+      param_.sizes(),
+      "; grad sizes: ",
+      grad_.sizes());
+  TORCH_CHECK(
+      param_.sizes() == exp_avg_.sizes(),
+      "Expect param and exp_avg have the same sizes, param sizes: ",
+      param_.sizes(),
+      "; exp_avg sizes: ",
+      exp_avg_.sizes());
+  TORCH_CHECK(
+      param_.sizes() == exp_avg_sq_.sizes(),
+      "Expect param and exp_avg_sq_ have the same sizes, param sizes: ",
+      param_.sizes(),
+      "; exp_avg_sq sizes: ",
+      exp_avg_sq_.sizes());
+  if (amsgrad) {
+    TORCH_CHECK(
+        param_.sizes() == max_exp_avg_sq_.sizes(),
+        "Expect param and max_exp_avg_sq_ have the same sizes, param sizes: ",
+        param_.sizes(),
+        "; max_exp_avg_sq sizes: ",
+        max_exp_avg_sq_.sizes());
+  }
+  TORCH_CHECK(
+      param2_.numel() == 0 || param_.sizes() == param2_.sizes(),
+      "Expect param and param2_ have the same sizes, param sizes: ",
+      param_.sizes(),
+      "; param2_ sizes: ",
+      param2_.sizes());
+
+  /*
+  pointer to adam_fused_step_kernel_impl(
+      param_,
+      exp_avg_,
+      exp_avg_sq_,
+      max_exp_avg_sq_,
+      grad_,
+      param2_,
+      amsgrad,
+      step,
+      beta1,
+      beta2,
+      learning_rate,
+      weight_decay,
+      eps);
+  */
+  adam_fused_step_kernel_stub(
+      kCPU,
+      param_,
+      exp_avg_,
+      exp_avg_sq_,
+      max_exp_avg_sq_,
+      grad_,
+      param2_,
+      amsgrad,
+      step,
+      beta1,
+      beta2,
+      learning_rate,
+      weight_decay,
+      eps);
+}
+
+} // namespace cpu
+} // namespace torch_ipex
+
+namespace {
+
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  m.def(
+      "adam_fused_step(Tensor(a!) param, Tensor(b!) exp_avg, Tensor(c!) "
+      "exp_avg_sq, Tensor(d!) max_exp_avg_sq, Tensor grad, Tensor trail, "
+      "bool amsgrad, float step, float beta1, float "
+      "beta2, float lr, float weight_decay, float eps) -> ()",
+      torch_ipex::cpu::adam_fused_step);
+}
+
+} // namespace

intel_extension_for_pytorch/csrc/aten/cpu/optimizer/AdamFusedStep.cpp

@@ -48,6 +48,21 @@ void packed_add_kernel_impl(
     const at::Tensor& grad,
     double alpha);
 
+void adam_fused_step_kernel_impl(
+    const at::Tensor& param_,
+    const at::Tensor& exp_avg_,
+    const at::Tensor& exp_avg_sq_,
+    const at::Tensor& max_exp_avg_sq_,
+    const at::Tensor& grad_,
+    const at::Tensor& param2_,
+    bool amsgrad,
+    double step,
+    double beta1,
+    double beta2,
+    double learning_rate,
+    double weight_decay,
+    double eps);
+
 } // namespace
 
 using adagrad_fused_step_kernel_fn = std::tuple<at::Tensor, at::Tensor> (*)(
@@ -93,5 +108,21 @@ using packed_add_kernel_fn =
     void (*)(at::Tensor&, at::Tensor&, const at::Tensor&, double);
 DECLARE_DISPATCH(packed_add_kernel_fn, packed_add_kernel_stub);
 
+using adam_fused_step_kernel_fn = void (*)(
+    const at::Tensor&,
+    const at::Tensor&,
+    const at::Tensor&,
+    const at::Tensor&,
+    const at::Tensor&,
+    const at::Tensor&,
+    bool,
+    double,
+    double,
+    double,
+    double,
+    double,
+    double);
+DECLARE_DISPATCH(adam_fused_step_kernel_fn, adam_fused_step_kernel_stub);
+
 } // namespace cpu
 } // namespace torch_ipex