support punica bgmv (#3422)

* enable punica

* change API

* enable bgmv_expand

* improve code style

* add kernel punica_bgmv_expand_slice

* add more ut; pass all ut

* add frontend APIs and docs; refine code

* revert APIs changes

* aligned APIs with xpu

* add docs

* add docs

* fix clang

Author: shiyang-weng

csrc/cpu/aten/Punica.cpp

@@ -0,0 +1,78 @@
+#include "Punica.h"
+#include <torch/all.h>
+#include <torch/csrc/autograd/function.h>
+#include "csrc/utils/CustomOperatorRegistration.h"
+
+namespace torch_ipex {
+namespace cpu {
+
+IPEX_DEFINE_DISPATCH(punica_bgmv_shrink_kernel_stub);
+IPEX_DEFINE_DISPATCH(punica_bgmv_expand_kernel_stub);
+IPEX_DEFINE_DISPATCH(punica_bgmv_expand_slice_kernel_stub);
+
+void punica_bgmv_shrink_forward_cpu(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    const double scale) {
+  return punica_bgmv_shrink_kernel_stub(
+      kCPU, out, input, weights, indicies, scale);
+}
+
+void punica_bgmv_expand_forward_cpu(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    bool add_inputs) {
+  return punica_bgmv_expand_kernel_stub(
+      kCPU, out, input, weights, indicies, add_inputs);
+}
+
+void punica_bgmv_expand_slice_forward_cpu(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    int64_t slice_offset,
+    int64_t slice_size,
+    bool add_inputs) {
+  return punica_bgmv_expand_slice_kernel_stub(
+      kCPU,
+      out,
+      input,
+      weights,
+      indicies,
+      slice_offset,
+      slice_size,
+      add_inputs);
+}
+
+} // namespace cpu
+} // namespace torch_ipex
+
+namespace {
+
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  IPEX_OP_REGISTER_DISPATCH(
+      "punica_bgmv_shrink",
+      torch_ipex::cpu::punica_bgmv_shrink_forward_cpu,
+      c10::DispatchKey::CPU);
+}
+
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  IPEX_OP_REGISTER_DISPATCH(
+      "punica_bgmv_expand",
+      torch_ipex::cpu::punica_bgmv_expand_forward_cpu,
+      c10::DispatchKey::CPU);
+}
+
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  IPEX_OP_REGISTER_DISPATCH(
+      "punica_bgmv_expand_slice",
+      torch_ipex::cpu::punica_bgmv_expand_slice_forward_cpu,
+      c10::DispatchKey::CPU);
+}
+
+} // namespace

csrc/cpu/aten/Punica.h

@@ -0,0 +1,67 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <dyndisp/DispatchStub.h>
+
+namespace torch_ipex {
+namespace cpu {
+
+namespace {
+
+void punica_bgmv_shrink(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    const double scale);
+
+void punica_bgmv_expand(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    bool add_inputs);
+
+void punica_bgmv_expand_slice(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    int64_t slice_offset,
+    int64_t slice_size,
+    bool add_inputs);
+} // namespace
+
+using punica_bgmv_shrink_fn = void (*)(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    const double scale);
+
+using punica_bgmv_expand_fn = void (*)(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    bool add_inputs);
+
+using punica_bgmv_expand_slice_fn = void (*)(
+    at::Tensor& out,
+    at::Tensor& input,
+    at::Tensor& weights,
+    at::Tensor& indicies,
+    int64_t slice_offset,
+    int64_t slice_size,
+    bool add_inputs);
+
+IPEX_DECLARE_DISPATCH(punica_bgmv_shrink_fn, punica_bgmv_shrink_kernel_stub);
+
+IPEX_DECLARE_DISPATCH(punica_bgmv_expand_fn, punica_bgmv_expand_kernel_stub);
+
+IPEX_DECLARE_DISPATCH(
+    punica_bgmv_expand_slice_fn,
+    punica_bgmv_expand_slice_kernel_stub);
+
+} // namespace cpu
+} // namespace torch_ipex