[Prim][PIR] Sink bn (#59054)

cyber-pioneer · web-flow · commit 781f3bfedfac · 2023-11-20T18:33:18.000+08:00
* sink bn

* fix template

* fix code

* bn manual

* fix code

* polish log

* fix code

* remove unused code

* pause test case

* fix test case
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -19,6 +19,7 @@
 # come into effect in generated file pd_op.h
 # manual decomp interface declare are located in manual_op.h
 decomp_interface_declare_gen_op_list = [
+    "batch_norm",
     "mean",
     "squeeze",
     "add_n",
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op_decomp.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op_decomp.cc
@@ -29,5 +29,93 @@ namespace paddle {
 namespace dialect {
 using IntArray = paddle::experimental::IntArray;
 
+std::vector<std::vector<pir::OpResult>> BatchNormOp::Decomp(
+    pir::Operation* op) {
+  VLOG(4) << "Decomp call batch_norm's decomp interface begin";
+  BatchNormOp op_obj = op->dyn_cast<BatchNormOp>();
+  (void)op_obj;
+
+  FLAGS_tensor_operants_mode = "static";
+
+  VLOG(6) << "Decomp Prepare inputs of batch_norm";
+
+  Tensor x(std::make_shared<primitive::LazyTensor>(op_obj.x()));
+  Tensor mean(std::make_shared<primitive::LazyTensor>(op_obj.mean()));
+  Tensor variance(std::make_shared<primitive::LazyTensor>(op_obj.variance()));
+  paddle::optional<Tensor> scale;
+  if (!IsEmptyValue(op_obj.scale())) {
+    scale = paddle::make_optional<Tensor>(
+        Tensor(std::make_shared<primitive::LazyTensor>(op_obj.scale())));
+  }
+  paddle::optional<Tensor> bias;
+  if (!IsEmptyValue(op_obj.bias())) {
+    bias = paddle::make_optional<Tensor>(
+        Tensor(std::make_shared<primitive::LazyTensor>(op_obj.bias())));
+  }
+
+  VLOG(6) << "Decomp prepare attributes of batch_norm";
+  bool is_test = op->attribute("is_test").dyn_cast<pir::BoolAttribute>().data();
+  float momentum =
+      op->attribute("momentum").dyn_cast<pir::FloatAttribute>().data();
+  float epsilon =
+      op->attribute("epsilon").dyn_cast<pir::FloatAttribute>().data();
+  const std::string& data_layout =
+      op->attribute("data_layout").dyn_cast<pir::StrAttribute>().AsString();
+  bool use_global_stats =
+      op->attribute("use_global_stats").dyn_cast<pir::BoolAttribute>().data();
+  bool trainable_statistics = op->attribute("trainable_statistics")
+                                  .dyn_cast<pir::BoolAttribute>()
+                                  .data();
+
+  VLOG(6) << "Decomp call batch_norm's forward composite rule prepare";
+
+  auto org_res = op->results();
+  std::vector<std::vector<pir::OpResult>> res(org_res.size());
+
+  VLOG(6) << "Decomp call batch_norm's forward composite rule begin";
+
+  std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> op_res =
+      paddle::primitive::details::batch_norm_decomp<primitive::LazyTensor>(
+          x,
+          mean,
+          variance,
+          scale,
+          bias,
+          is_test,
+          momentum,
+          epsilon,
+          data_layout,
+          use_global_stats,
+          trainable_statistics);
+
+  VLOG(6) << "Decomp call batch_norm's forward composite rule end";
+
+  res[0].push_back(std::static_pointer_cast<primitive::LazyTensor>(
+                       std::get<0>(op_res).impl())
+                       ->value()
+                       .dyn_cast<pir::OpResult>());
+  res[1].push_back(std::static_pointer_cast<primitive::LazyTensor>(
+                       std::get<1>(op_res).impl())
+                       ->value()
+                       .dyn_cast<pir::OpResult>());
+  res[2].push_back(std::static_pointer_cast<primitive::LazyTensor>(
+                       std::get<2>(op_res).impl())
+                       ->value()
+                       .dyn_cast<pir::OpResult>());
+  res[3].push_back(std::static_pointer_cast<primitive::LazyTensor>(
+                       std::get<3>(op_res).impl())
+                       ->value()
+                       .dyn_cast<pir::OpResult>());
+  res[4].push_back(std::static_pointer_cast<primitive::LazyTensor>(
+                       std::get<4>(op_res).impl())
+                       ->value()
+                       .dyn_cast<pir::OpResult>());
+  pir::OpResult reserve_space;
+  res[5].push_back(reserve_space);
+
+  VLOG(4) << "Decomp call batch_norm's decomp interface end";
+  return res;
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/codegen/decomp_gen.py b/paddle/fluid/primitive/codegen/decomp_gen.py
@@ -197,7 +197,7 @@ def gen(
         for attr_item in item["attrs"]:
             if attr_item["typename"] not in attr_types_map.keys():
                 raise TypeError
-            attr_item["mapped_type"] = attr_types_map[attr_item["typename"]][0]
+            attr_item["mapped_type"] = attr_types_map[attr_item["typename"]]
         for out_item in item["outputs"]:
             if out_item["typename"] not in output_type_map.keys():
                 name = out_item["typename"]
diff --git a/paddle/fluid/primitive/codegen/templates/decomp/generated_decomp.j2 b/paddle/fluid/primitive/codegen/templates/decomp/generated_decomp.j2
@@ -21,12 +21,14 @@ using IntArray = paddle::experimental::IntArray;
 {% set output_types=[] %}
 
 std::vector<std::vector<pir::OpResult>> {{class_name}}::Decomp(pir::Operation* op) {
+  VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface begin";
+
   {{class_name}} op_obj = op->dyn_cast<{{class_name}}>();
   (void)op_obj;
 
   FLAGS_tensor_operants_mode = "static";
 
-  VLOG(4) << "Decomp Prepare inputs of {{fwd_name}}";
+  VLOG(6) << "Decomp Prepare inputs of {{fwd_name}}";
 
     {% for item in inputs -%}
       {% do input_names.append(item.name) %}
@@ -66,7 +68,7 @@ std::vector<std::vector<pir::OpResult>> {{class_name}}::Decomp(pir::Operation* o
       {% endif %}
     {% endfor %}
 
-  VLOG(4) << "Decomp prepare attributes of {{fwd_name}}";
+  VLOG(6) << "Decomp prepare attributes of {{fwd_name}}";
     {% if attrs %}
       {% for item in attrs %}
       {% do attr_names.append(item.name) %}
@@ -106,18 +108,25 @@ std::vector<std::vector<pir::OpResult>> {{class_name}}::Decomp(pir::Operation* o
       paddle::dialect::GetInt64Vector({{item.name}}_define_op->attribute("value")));
 
       {% else %}
-  {{item.typename}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type}}>().data();
+        {% if item.mapped_type[0] == "pir::StrAttribute" %}
+  {{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().AsString();
+        {% else %}
+  {{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().data();
+        {% endif %}
       {% endif %}
       {% endfor %}
     {% endif %}
 
-  VLOG(4) << "Decomp prepare call {{fwd_name}}'s decomp interface";
+  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule prepare";
 
   auto org_res = op->results();
   std::vector<std::vector<pir::OpResult>> res(org_res.size());
 
+  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule begin";
   {% if outputs|length == 1 %}
   Tensor op_res = paddle::primitive::details::{{fwd_name}}_decomp<primitive::LazyTensor>({{common.args(input_names, attr_names)}});
+  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule end";
+
   res[0].push_back(
     std::static_pointer_cast<primitive::LazyTensor>(op_res.impl())
         ->value()
@@ -129,6 +138,7 @@ std::vector<std::vector<pir::OpResult>> {{class_name}}::Decomp(pir::Operation* o
     {% endfor %}
   std::tuple<{{common.sequence('', '', ', ', output_types)}}> op_res = paddle::primitive::details::{{fwd_name}}_decomp<primitive::LazyTensor>(
           {{common.args(input_names, attr_names)}});
+  VLOG(6) << "Decomp call {{fwd_name}}'s forward composite rule end";
   {% for k in range(outputs|length) %}
     {% if outputs[k].intermediate and fwd_name in decomp_ops_list_contain_unused_output %}
   pir::OpResult {{outputs[k].name}};
@@ -139,6 +149,7 @@ std::vector<std::vector<pir::OpResult>> {{class_name}}::Decomp(pir::Operation* o
   {% endfor %}
   {% endif %}
 
+  VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface end";
   return res;
 }
 {% endmacro %}
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -62,6 +62,122 @@ Tensor mean_decomp(const Tensor& x, const IntArray& axis, bool keepdim) {
   }
 }
 
+template <typename T>
+std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> batch_norm_decomp(
+    const Tensor& x,
+    const Tensor& run_mean,
+    const Tensor& run_var,
+    const paddle::optional<Tensor>& scale,
+    const paddle::optional<Tensor>& bias,
+    bool is_test,
+    float momentum,
+    float epsilon,
+    const std::string& data_layout,
+    bool use_global_stats,
+    bool trainable_statistics) {
+  auto org_dtype = x.dtype();
+  Tensor x_cast = x;
+
+  bool need_cast = org_dtype == phi::DataType::FLOAT16 ||
+                   org_dtype == phi::DataType::BFLOAT16;
+
+  if (need_cast) {
+    x_cast = cast<T>(x, phi::DataType::FLOAT32);
+  }
+
+  std::vector<int64_t> x_dim = phi::vectorize<int64_t>(x_cast.dims());
+  int rank = x_dim.size();
+  DataLayout data_layout_ = phi::StringToDataLayout(data_layout);
+  int feature_axis;
+  if (data_layout_ == DataLayout::kNCHW) {
+    feature_axis = 1;
+  } else if (data_layout_ == DataLayout::kNHWC) {
+    feature_axis = rank - 1;
+  } else {
+    PADDLE_THROW(
+        phi::errors::InvalidArgument("Unknown storage order: %s", data_layout));
+  }
+  std::vector<int64_t> reduce_axes;
+  for (int i = 0; i < rank; ++i) {
+    if (i != feature_axis) {
+      reduce_axes.push_back(i);
+    }
+  }
+  std::vector<int64_t> stats_shape;
+  for (int i = 0; i < rank; ++i) {
+    if (find_value(reduce_axes, i) == false) {
+      stats_shape.push_back(x_dim[i]);
+    } else {
+      stats_shape.push_back(1);
+    }
+  }
+
+  Tensor half = full<T>(IntArray({1}), -0.5, x_cast.dtype());
+
+  bool use_run_stat = (is_test && (!trainable_statistics)) || use_global_stats;
+  Tensor x_hat;
+  Tensor batch_mean;
+  Tensor inv_std;
+  Tensor run_mean_;
+  Tensor run_var_;
+  if (!use_run_stat) {
+    batch_mean = mean_decomp<T>(x_cast, IntArray(reduce_axes), false);
+    auto temp = mean_decomp<T>(x_cast * x_cast, IntArray(reduce_axes), false);
+    auto batch_var = temp - batch_mean * batch_mean;
+    inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+    if (data_layout_ == DataLayout::kNHWC) {
+      x_hat = (x_cast - batch_mean) * inv_std;
+    } else {
+      x_hat = (x_cast - reshape<T>(batch_mean, stats_shape)) *
+              reshape<T>(inv_std, stats_shape);
+    }
+    run_mean_ = run_mean * momentum + batch_mean * (1. - momentum);
+    run_var_ = run_var * momentum + batch_var * (1. - momentum);
+  } else {
+    batch_mean = full<T>(phi::vectorize(run_mean.dims()), 0, run_mean.dtype());
+    auto batch_var =
+        full<T>(phi::vectorize(run_var.dims()), 0, run_var.dtype());
+    inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+    if (data_layout_ == DataLayout::kNHWC) {
+      x_hat =
+          (x_cast - run_mean) * elementwise_pow<T>((run_var + epsilon), half);
+    } else {
+      x_hat = (x_cast - reshape<T>(run_mean, stats_shape)) *
+              elementwise_pow<T>((reshape<T>(run_var, stats_shape) + epsilon),
+                                 half);
+    }
+    run_mean_ = assign<T>(run_mean);
+    run_var_ = assign<T>(run_var);
+  }
+  Tensor y;
+  Tensor new_scale =
+      scale ? scale.get()
+            : full<T>(phi::vectorize(x_cast.dims()), 1, x_cast.dtype());
+  Tensor new_bias =
+      bias ? bias.get()
+           : full<T>(phi::vectorize(x_cast.dims()), 0, x_cast.dtype());
+  if (data_layout_ == DataLayout::kNHWC) {
+    y = x_hat * new_scale + new_bias;
+  } else {
+    y = x_hat * reshape<T>(new_scale, stats_shape) +
+        reshape<T>(new_bias, stats_shape);
+  }
+  if (need_cast) {
+    y = cast<T>(y, org_dtype);
+  }
+  Tensor reserve_space;
+
+  auto batch_mean_ = assign<T>(batch_mean);
+  auto inv_std_ = assign<T>(inv_std);
+  if (!use_run_stat) {
+    return std::make_tuple(
+        y, run_mean_, run_var_, batch_mean_, inv_std_, reserve_space);
+  } else {
+    return std::make_tuple(
+        y, run_mean_, run_var_, batch_mean_, inv_std_, reserve_space);
+  }
+}
+
 template <typename T>
 Tensor softmax_decomp(const Tensor& x, const int& axis) {
   auto org_dtype = x.dtype();
diff --git a/paddle/fluid/primitive/utils/utils.h b/paddle/fluid/primitive/utils/utils.h
@@ -139,5 +139,13 @@ std::vector<std::vector<Tensor>> ConstructVjpResultByStopGradients(
     const std::vector<std::vector<Tensor>>& outputs,
     const std::vector<std::vector<bool>>& stop_gradients);
 
+static bool find_value(const std::vector<int64_t>& vec, int64_t value) {
+  if (std::find(vec.begin(), vec.end(), value) != vec.end()) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/base/core.py b/python/paddle/base/core.py
@@ -503,8 +503,14 @@ def _get_batch_norm_none_var(op):
     "unsqueeze2": ["XShape"],
 }
 
+
 # some intermediate outputs like xshape will no longer used after decomp, but return none to keep output num the same as origin op
-decomp_ops_list_contain_unused_output = ["pd_op.squeeze", "pd_op.unsqueeze"]
+# key is the name of op, and value is the index of output in op.outputs
+decomp_ops_contain_unused_output = {
+    "pd_op.squeeze": [1],
+    "pd_op.unsqueeze": [1],
+    "pd_op.batch_norm": [5],
+}
 
 
 def _set_prim_forward_blacklist(*args):
diff --git a/python/paddle/decomposition/decomp.py b/python/paddle/decomposition/decomp.py
@@ -19,7 +19,7 @@
 from paddle.autograd import ir_backward
 from paddle.base.core import (
     call_decomp,
-    decomp_ops_list_contain_unused_output,
+    decomp_ops_contain_unused_output,
     has_decomp,
 )
 from paddle.base.libpaddle.pir import Block, Operation, Program
@@ -37,22 +37,20 @@ def _build_tensor_tuple(xs):
 
 
 def _analyse_decomp_results(orig_outs, decomp_outs, op):
-    intermediate_status = op.get_output_intermediate_status()
-    assert len(orig_outs) == len(decomp_outs) == len(intermediate_status)
+    assert len(orig_outs) == len(decomp_outs)
     res = []
-    for org_item, new_item, value in zip(
-        orig_outs, decomp_outs, intermediate_status
-    ):
-        if isinstance(org_item, pir.OpResult):
-            if value and op.name() in decomp_ops_list_contain_unused_output:
-                assert new_item[0] is None
+    for idx, value in enumerate(decomp_outs):
+        if isinstance(orig_outs[idx], pir.OpResult):
+            if (
+                op.name() in decomp_ops_contain_unused_output.keys()
+                and idx in decomp_ops_contain_unused_output[op.name()]
+            ):
+                assert value[0] is None
             else:
-                assert len(new_item) == 1 and isinstance(
-                    new_item[0], pir.OpResult
-                )
-            res.append(new_item[0])
+                assert len(value) == 1 and isinstance(value[0], pir.OpResult)
+            res.append(value[0])
         else:
-            res.append(new_item)
+            res.append(value)
     return res
 
 
@@ -278,10 +276,18 @@ def _decompose_subgraph(block, orig_vars, dst_vars, op_filter):
                 _check_op_results(
                     op_name, orig_outs, new_outs, orig_vars, dst_vars
                 )
-                if op.name() in decomp_ops_list_contain_unused_output:
-                    orig_outs[0].replace_all_uses_with(new_outs[0])
+                if op.name() in decomp_ops_contain_unused_output.keys():
+                    for idx in range(len(orig_outs)):
+                        if (
+                            idx
+                            not in decomp_ops_contain_unused_output[op.name()]
+                        ):
+                            orig_outs[idx].replace_all_uses_with(new_outs[idx])
                 else:
-                    op.replace_all_uses_with(new_outs)
+                    if op.name() in decomp_ops_contain_unused_output.keys():
+                        orig_outs[0].replace_all_uses_with(new_outs[0])
+                    else:
+                        op.replace_all_uses_with(new_outs)
                 block.remove_op(op)
 
                 if temp_op is not None:
diff --git a/test/legacy_test/test_batch_norm_op_prim_nchw.py b/test/legacy_test/test_batch_norm_op_prim_nchw.py
diff --git a/test/legacy_test/test_batch_norm_op_prim_nhwc.py b/test/legacy_test/test_batch_norm_op_prim_nhwc.py
