[Prim][PIR] support instancenorm op dynamic forward in prim pir (#64598)

* support dynamic instancenorm

* fix narrow convert

* fix dtype

* update

* remove empty_shape

* Update composite.h

* Update composite.h

* fix scalar

Author: Eddie-Wang1120

paddle/fluid/primitive/base/decomp_trans.cc

@@ -46,15 +46,13 @@ std::unordered_set<std::string> decomp_op_contain_none = {"pd_op.squeeze",
                                                           "pd_op.batch_norm_",
                                                           "pd_op.dropout"};
 //
-std::unordered_set<std::string> dynamic_shape_blacklist = {
-    "pd_op.squeeze",
-    "pd_op.unsqueeze",
-    "pd_op.batch_norm",
-    "pd_op.batch_norm_",
-    "pd_op.bmm",
-    "pd_op.flatten",
-    "pd_op.instance_norm",
-    "pd_op.one_hot"};
+std::unordered_set<std::string> dynamic_shape_blacklist = {"pd_op.squeeze",
+                                                           "pd_op.unsqueeze",
+                                                           "pd_op.batch_norm",
+                                                           "pd_op.batch_norm_",
+                                                           "pd_op.bmm",
+                                                           "pd_op.flatten",
+                                                           "pd_op.one_hot"};
 
 namespace {
 std::set<std::string> StringSplit(const std::string& str) {

paddle/fluid/primitive/composite/composite.h

@@ -953,6 +953,78 @@ std::tuple<Tensor, Tensor, Tensor> instance_norm_decomp(
     const paddle::optional<Tensor>& scale,
     const paddle::optional<Tensor>& bias,
     float epsilon) {
+  if (has_dynamic_shape(x.shape())) {
+    auto org_dtype = x.dtype();
+    Tensor x_cast = x;
+
+    bool need_cast = is_half_dtype(org_dtype);
+    if (need_cast) {
+      x_cast = cast<T>(x, DataType::FLOAT32);
+    }
+
+    std::vector<int64_t> axis;
+    auto x_dim = x.shape();
+    for (size_t i = 2; i < x_dim.size(); i++) {
+      axis.push_back(static_cast<int64_t>(i));
+    }
+
+    // out = (x - mean(x)) / sqrt(var + epsilon))
+    // var = mean((x-mean(x))^2)
+    auto mean_ = mean_decomp<T>(x_cast, axis, true);
+    auto difference = x_cast - mean_;
+    auto var_tmp1 = difference * difference;
+    auto variance = mean_decomp<T>(var_tmp1, axis, true);
+    auto var_shape = shape<T>(variance);
+    auto var_tmp3 = variance + full_scalar<T>(epsilon, variance.dtype());
+    auto rsqrt_var = rsqrt<T>(var_tmp3);
+    auto out = difference * rsqrt_var;
+
+    int dim_size = x_dim.size();
+    auto x_shape_tensor = shape<T>(x);
+    std::vector<Tensor> slice_shape_concat;
+
+    auto shape_1 = full<T>({1}, 1, x_shape_tensor.dtype());
+    auto shape_2 =
+        cast<T>(get_slice<T>(x_shape_tensor, 1), x_shape_tensor.dtype());
+    auto shape_3 = full<T>({dim_size - 2}, 1, x_shape_tensor.dtype());
+
+    slice_shape_concat.push_back(shape_1);
+    slice_shape_concat.push_back(shape_2);
+    slice_shape_concat.push_back(shape_3);
+    auto slice_shape_tensor = concat<T>(slice_shape_concat, 0);
+
+    Tensor scale_cast;
+    if (scale) {
+      scale_cast =
+          backend::reshape_with_tensor<T>(scale.get(), slice_shape_tensor);
+      if (need_cast) {
+        scale_cast = cast<T>(scale_cast, DataType::FLOAT32);
+      }
+      out = out * scale_cast;
+    }
+    Tensor bias_cast;
+    if (bias) {
+      bias_cast =
+          backend::reshape_with_tensor<T>(bias.get(), slice_shape_tensor);
+      if (need_cast) {
+        bias_cast = cast<T>(bias_cast, DataType::FLOAT32);
+      }
+      out = out + bias_cast;
+    }
+
+    std::vector<int64_t> res_shape(1, -1);
+    auto mean_out = reshape<T>(mean_, res_shape);
+    auto variance_out = reshape<T>(rsqrt_var, res_shape);
+
+    Tensor res;
+    if (need_cast) {
+      res = cast<T>(out, org_dtype);
+    } else {
+      res = out;
+    }
+
+    return std::make_tuple(res, mean_out, variance_out);
+  }
   auto org_dtype = x.dtype();
   Tensor x_cast = x;
 

test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py

@@ -164,6 +164,10 @@ def layer_norm_net1(x):
     return paddle.nn.functional.layer_norm(x, x.shape[1:])
 
 
+def instance_norm_net(x):
+    return paddle.nn.functional.instance_norm(x)
+
+
 def flatten_net(x):
     return paddle.flatten(x, 1, 2)
 
@@ -488,6 +492,19 @@ def setUp(self):
         self.tol = 5e-6
 
 
+class TestPrimInstancenorm(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2, 32, 128]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = instance_norm_net
+        self.necessary_ops = "pd_op.instance_norm"
+        self.enable_cinn = False
+        self.tol = 5e-6
+
+
 class TestPrimGroupNorm1(TestPrimBase):
     def setUp(self):
         np.random.seed(2023)