add linear_swish fusion using MKL for FSI riskfuel after code review. (#551)

tangleintel · EikanWang · web-flow · commit cc855ff2bafd · 2022-02-23T17:18:09.000+08:00
* add linear_swish fusion using MKL for FSI riskfuel
Co-authored-by: Wang Weihan &lt;eikan.wang@intel.com&gt;
diff --git a/intel_extension_for_pytorch/csrc/aten/cpu/kernels/jit_kernels/AddSwishKrnl.cpp b/intel_extension_for_pytorch/csrc/aten/cpu/kernels/jit_kernels/AddSwishKrnl.cpp
@@ -0,0 +1,41 @@
+#include <csrc/jit/cpu/kernels/AddSwish.h>
+
+#if defined(CPU_CAPABILITY_AVX512)
+#include "csrc/cpu/vec512/add_swish.h"
+#endif
+
+namespace torch_ipex {
+namespace cpu {
+
+#if defined(DYN_DISP_BUILD)
+namespace {
+#endif
+
+at::Tensor add_swish_kernel_impl(
+    at::Tensor& x,
+    at::Tensor& a,
+    const at::Tensor& b,
+    const at::Tensor& c) {
+#if defined(CPU_CAPABILITY_AVX512)
+  if (a.scalar_type() == at::kFloat && c.scalar_type() == at::kFloat) {
+    return torch_ipex::cpu::kernel::vec::vec512::dil_add_swish<float>(a, c);
+  } else if (
+      a.scalar_type() == at::kBFloat16 && c.scalar_type() == at::kBFloat16) {
+    return torch_ipex::cpu::kernel::vec::vec512::dil_add_swish<at::BFloat16>(
+        a, c);
+  }
+#endif
+  auto lin_res = at::linear(x, b, c);
+  auto sigmoid_res = at::sigmoid(lin_res);
+  return at::mul(lin_res, sigmoid_res);
+}
+
+#if defined(DYN_DISP_BUILD)
+} // anonymous namespace
+
+REGISTER_DISPATCH(add_swish_kernel_stub, &add_swish_kernel_impl);
+
+#endif
+
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/cpu/vec512/add_swish.h b/intel_extension_for_pytorch/csrc/cpu/vec512/add_swish.h
@@ -0,0 +1,109 @@
+#pragma once
+
+#include <immintrin.h>
+
+#include <ATen/ATen.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/Parallel.h>
+#include <c10/util/SmallVector.h>
+#include <limits>
+#include "add_softmax.h"
+#include "utils.h"
+
+namespace torch_ipex {
+namespace cpu {
+namespace kernel {
+namespace vec {
+namespace vec512 {
+
+template <typename scalar_t>
+inline void _dil_add_swish_fusion_kernel(
+    scalar_t* a,
+    const scalar_t* b,
+    const int& size) {
+  auto vec_ps_min = _mm512_set1_ps(std::numeric_limits<float>::min());
+  auto vec_ps_1 = _mm512_set1_ps(1.0);
+  __m512 vec_a, vec_b;
+  __m512 vec_add_tmp, vec_addone_tmp;
+
+  int i = 0;
+
+  // load tensor<float> a & b
+  // assum the same size , no need to broadcast
+  for (; i <= size - 16; i += 16) {
+    // a is first operand of add, b is bias
+    vec_a = _loadu(a + i);
+    vec_b = _loadu(b + i);
+
+    // add bias
+    vec_a = _mm512_add_ps(vec_a, vec_b);
+    vec_add_tmp =
+        vec_a; // keep the intermediate result for later use in the mul
+
+    // caculate sigmoid e^x / (1 + e^x)
+    vec_a = _dil_exp_kernel(vec_a);
+    vec_addone_tmp = _mm512_add_ps(vec_a, vec_ps_1);
+    vec_a = _mm512_div_ps(vec_a, vec_addone_tmp);
+    vec_a = _mm512_mul_ps(vec_a, vec_add_tmp);
+
+    _storeu(a + i, vec_a);
+  }
+
+  // 512 tail
+  if (i < size) {
+    // mask load
+    __mmask16 mask = (1 << (size - i)) - 1;
+    vec_a = _maskz_loadu(a + i, mask);
+    vec_b = _maskz_loadu(b + i, mask);
+
+    // add bias
+    vec_a = _mm512_add_ps(vec_a, vec_b);
+    vec_add_tmp =
+        vec_a; // keep the intermediate result for later use in the second mul
+
+    // caculate sigmoid e^x / (1 + e^x)
+    vec_a = _dil_exp_kernel(vec_a);
+    vec_addone_tmp = _mm512_add_ps(vec_a, vec_ps_1);
+    vec_a = _mm512_div_ps(vec_a, vec_addone_tmp);
+
+    vec_a = _mm512_mul_ps(vec_a, vec_add_tmp);
+
+    // mask store
+    _mask_storeu(a + i, vec_a, mask);
+  }
+}
+
+template <typename scalar_t>
+at::Tensor dil_add_swish(const at::Tensor& mm_output, const at::Tensor& bias) {
+  scalar_t* mm_output_data_base = mm_output.data_ptr<scalar_t>();
+  scalar_t* bias_data_base = bias.data_ptr<scalar_t>();
+
+  auto infered_size = mm_output.sizes().vec();
+  int64_t dim_size = infered_size[infered_size.size() - 1];
+  int64_t outer_size = 1;
+  // The last dim is the loop unit. We need to minus 2 to exclude the last dim.
+  // infered_size.size() - 2 is the -2th dimension.
+  for (int64_t i = infered_size.size() - 2; i >= 0; i--) {
+    // Calculate outer loop number;
+    outer_size *= infered_size[i];
+  }
+
+  int64_t grain_size = at::internal::GRAIN_SIZE / (16 * dim_size);
+  if (grain_size < 1)
+    grain_size = 1;
+
+  at::parallel_for(0, outer_size, grain_size, [&](int64_t begin, int64_t end) {
+    for (int64_t i = begin; i < end; i++) {
+      _dil_add_swish_fusion_kernel<scalar_t>(
+          mm_output_data_base + i * dim_size, bias_data_base, dim_size);
+    }
+  });
+
+  return mm_output;
+} // dil_add_swish
+
+} // namespace vec512
+} // namespace vec
+} // namespace kernel
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/kernels/AddSwish.cpp b/intel_extension_for_pytorch/csrc/jit/cpu/kernels/AddSwish.cpp
@@ -0,0 +1,29 @@
+#include "AddSwish.h"
+#include <ATen/Context.h>
+#include <ATen/InferSize.h>
+#include <c10/util/Exception.h>
+#include <c10/util/Logging.h>
+#include <torch/csrc/autograd/function.h>
+#include <iostream>
+
+#include <limits>
+
+namespace torch_ipex {
+namespace cpu {
+DEFINE_DISPATCH(add_swish_kernel_stub);
+
+// Currently we only support 1D tensor of bias(operand of add).
+at::Tensor AddSwish(
+    at::Tensor& x,
+    at::Tensor& mm_output,
+    const at::Tensor& weight,
+    const at::Tensor& bias) {
+#if defined(DYN_DISP_BUILD)
+  return add_swish_kernel_stub(kCPU, x, mm_output, weight, bias);
+#else
+  return add_swish_kernel_impl(x, mm_output, weight, bias);
+#endif
+}
+
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/kernels/AddSwish.h b/intel_extension_for_pytorch/csrc/jit/cpu/kernels/AddSwish.h
@@ -0,0 +1,38 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <csrc/dyndisp/DispatchStub.h>
+
+namespace torch_ipex {
+namespace cpu {
+
+// Currently we only support 1D tensor of bias(operand of add).
+at::Tensor AddSwish(
+    at::Tensor& x,
+    at::Tensor& mm_output,
+    const at::Tensor& weight,
+    const at::Tensor& bias);
+
+#if defined(DYN_DISP_BUILD)
+namespace {
+#endif
+
+at::Tensor add_swish_kernel_impl(
+    at::Tensor& x,
+    at::Tensor& a,
+    const at::Tensor& b,
+    const at::Tensor& c);
+
+#if defined(DYN_DISP_BUILD)
+}
+#endif
+
+using add_swish_kernel_fn = at::Tensor (*)(
+    at::Tensor&,
+    at::Tensor&,
+    const at::Tensor&,
+    const at::Tensor&);
+DECLARE_DISPATCH(add_swish_kernel_fn, add_swish_kernel_stub);
+
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/kernels/LinearSwishCustomized.cpp b/intel_extension_for_pytorch/csrc/jit/cpu/kernels/LinearSwishCustomized.cpp
@@ -0,0 +1,32 @@
+#include "LinearSwishCustomized.h"
+#include "AddSwish.h"
+
+#include <ATen/Context.h>
+#include <ATen/InferSize.h>
+#include <c10/util/Exception.h>
+#include <c10/util/Logging.h>
+#include <torch/csrc/autograd/function.h>
+
+#include <limits>
+
+#include "csrc/cpu/ideep/ideep.hpp"
+#include "csrc/utils/ipex_op_profile.h"
+
+namespace torch_ipex {
+namespace cpu {
+
+at::Tensor dil_linear_swish_customized(
+    at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& bias) {
+#if defined(IPEX_PROFILE_OP)
+  RECORD_FUNCTION("dil_linear_swish_customized", std::vector<c10::IValue>({}));
+#endif
+
+  // at::linear w/o bias
+  auto linear_res = at::linear(x, weight);
+  return AddSwish(x, linear_res, weight, bias);
+}
+
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/kernels/LinearSwishCustomized.h b/intel_extension_for_pytorch/csrc/jit/cpu/kernels/LinearSwishCustomized.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include <ATen/Tensor.h>
+
+#include <c10/core/Scalar.h>
+#include <torch/csrc/jit/runtime/custom_operator.h>
+
+#include "csrc/cpu/ideep/ideep.hpp"
+
+namespace torch_ipex {
+namespace cpu {
+
+at::Tensor dil_linear_swish_customized(
+    at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& bias);
+} // namespace cpu
+} // namespace torch_ipex
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/passes/graph_rewrite.cpp b/intel_extension_for_pytorch/csrc/jit/cpu/passes/graph_rewrite.cpp
@@ -513,6 +513,24 @@ void FuseConcatBnRelu(std::shared_ptr<Graph>& graph) {
   rewriter_concatbnrelu.runOnGraph(graph, fusion_filter);
 }
 
+void FuseLinearSwishCustomized(std::shared_ptr<Graph>& graph) {
+  std::string linear_swish = R"(
+      graph(%x, %weight, %bias):
+        %_linear_res = aten::linear(%x, %weight, %bias)
+        %_sigmod_res = aten::sigmoid(%_linear_res)
+        %_mul_res2 = aten::mul(%_linear_res, %_sigmod_res)
+        return (%_mul_res2) )";
+
+  std::string linear_swish_fusion = R"(
+      graph(%x, %weight, %bias):
+        %_res = ipex::linear_swish_customized(%x, %weight, %bias)
+        return (%_res) )";
+
+  SubgraphRewriter ls_fusion;
+  ls_fusion.RegisterRewritePattern(linear_swish, linear_swish_fusion);
+  ls_fusion.runOnGraph(graph);
+}
+
 } // namespace graph_rewrite
 } // namespace jit
 } // namespace torch
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/passes/graph_rewrite.h b/intel_extension_for_pytorch/csrc/jit/cpu/passes/graph_rewrite.h
@@ -25,6 +25,7 @@ c10::optional<IValue> getIValue(
 
 void FuseShuffle(std::shared_ptr<Graph>& graph);
 void FuseMHAScoreCalc(std::shared_ptr<Graph>& graph);
+void FuseLinearSwishCustomized(std::shared_ptr<Graph>& graph);
 void replaceAtenMaxPool2dWithIpexMaxPool2d(std::shared_ptr<Graph>& graph);
 void replaceOpsWithAtenInplaceOps(std::shared_ptr<Graph>& graph);
 void replaceAtenOpsWithIpexInplaceOps(std::shared_ptr<Graph>& graph);
diff --git a/intel_extension_for_pytorch/csrc/jit/cpu/passes/register_dnnl_jit_ops.cpp b/intel_extension_for_pytorch/csrc/jit/cpu/passes/register_dnnl_jit_ops.cpp
@@ -10,6 +10,7 @@
 #include "csrc/jit/cpu/kernels/Embeddingbag.h"
 #include "csrc/jit/cpu/kernels/Interaction.h"
 #include "csrc/jit/cpu/kernels/LinearPacked.h"
+#include "csrc/jit/cpu/kernels/LinearSwishCustomized.h"
 #include "csrc/jit/cpu/kernels/Matmul.h"
 #include "csrc/jit/cpu/kernels/MaxPool2D.h"
 #include "csrc/jit/cpu/kernels/Mha.h"
@@ -521,6 +522,18 @@ RegisterOperators op({
         },
         aliasAnalysisFromSchema()),
 
+    Operator(
+        "ipex::linear_swish_customized(Tensor x, Tensor weight, Tensor ? bias) -> Tensor",
+        [](Stack& stack) {
+          auto result = dil_linear_swish_customized(
+              peek(stack, 0, 3).toTensor(),
+              peek(stack, 1, 3).toTensor(),
+              toOptionalTensor(std::move(peek(stack, 2, 3))));
+          drop(stack, 3);
+          pack(stack, std::move(result));
+        },
+        aliasAnalysisFromSchema()),
+
     Operator(
         "ipex::distil_mha_scores_calc(Tensor q, Tensor k, Tensor mask_qk, "
         "int[] mask_qk_reshp, int transpose_dim_a, int transpose_dim_b, "
@@ -539,6 +552,20 @@ RegisterOperators op({
               peek(stack, 8, 10).toInt(),
               peek(stack, 9, 10));
           drop(stack, 10);
+
+          pack(stack, std::move(result));
+        },
+        aliasAnalysisFromSchema()),
+
+    Operator(
+        "ipex::linear_swish_customized(Tensor x, Tensor weight, Tensor ? bias) -> Tensor",
+        [](Stack& stack) {
+          auto result = dil_linear_swish_customized(
+              peek(stack, 0, 3).toTensor(),
+              peek(stack, 1, 3).toTensor(),
+              toOptionalTensor(std::move(peek(stack, 2, 3))));
+          drop(stack, 3);
+
           pack(stack, std::move(result));
         },
         aliasAnalysisFromSchema()),
diff --git a/intel_extension_for_pytorch/csrc/jit/fusion_pass.cpp b/intel_extension_for_pytorch/csrc/jit/fusion_pass.cpp
@@ -377,6 +377,7 @@ void IPEXFusionPass(std::shared_ptr<Graph>& graph) {
   graph_rewrite::fuseLinearWithEltwise(graph);
   graph_rewrite::fuseLinearAddRelu(graph);
 
+  graph_rewrite::FuseLinearSwishCustomized(graph);
   // fuse add+layernorm
   graph_rewrite::FuseAddLayerNorm(graph);
   // deconvolution fusion
diff --git a/tests/cpu/test_jit.py b/tests/cpu/test_jit.py
