[Prim][PIR] Support the dynamic shape for one_hot op (#68295)

zeroRains · web-flow · commit dabf4221d996 · 2024-09-20T09:50:10.000+08:00
* suport the dynamic shape for one_hot op

* remove the cout

* fix the bug in cout
diff --git a/paddle/fluid/primitive/base/decomp_trans.cc b/paddle/fluid/primitive/base/decomp_trans.cc
@@ -54,8 +54,7 @@ std::unordered_set<std::string> dynamic_shape_blacklist = {"pd_op.squeeze",
                                                            "pd_op.batch_norm",
                                                            "pd_op.batch_norm_",
                                                            "pd_op.bmm",
-                                                           "pd_op.flatten",
-                                                           "pd_op.one_hot"};
+                                                           "pd_op.flatten"};
 
 namespace {
 std::set<std::string> StringSplit(const std::string& str) {
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -211,44 +211,72 @@ template <typename T>
 Tensor one_hot_decomp(const Tensor& x, const Tensor& num_classes) {
   auto num_classes_tensor =
       backend::full_with_tensor<T>(num_classes, 0, x.dtype());
+  Tensor ans;
+  if (has_dynamic_shape(x.shape())) {
+    Tensor x_shape = shape<T>(x);
+    Tensor one_tensor = full<T>({1}, 1, x_shape.dtype());
+    Tensor zero_tensor = full<T>({1}, 0, x_shape.dtype());
+    Tensor x_dims = one_tensor;
+    for (size_t i = 0; i < x.shape().size(); i++) {
+      x_dims = x_dims * get_slice<T>(x_shape, i);
+    }
+    Tensor input_dim = concat<T>(
+        {x_dims, full<T>({1}, num_classes_tensor.shape()[0], x_shape.dtype())});
+    auto input_tensor = backend::full_with_tensor<T>(input_dim, 0, x.dtype());
+    auto output_dim = concat<T>(
+        {x_shape,
+         full<T>({1}, num_classes_tensor.shape()[0], x_shape.dtype())});
+    auto arange_tensor = backend::arange_with_tensor<T>(
+        zero_tensor, x_dims, one_tensor, x_shape.dtype());
+    Tensor reshape_dim = concat<T>({x_dims, one_tensor});
+    auto x_reshape = backend::reshape<T>(x, reshape_dim);
+    auto arange_tensor_reshape =
+        backend::reshape<T>(arange_tensor, reshape_dim);
+    auto index_tensor = concat<T>({arange_tensor_reshape, x_reshape}, 1);
+    auto update_tensor = backend::full_with_tensor<T>(x_dims, 1, x.dtype());
+    ans = backend::reshape<T>(
+        cast<T>(scatter_nd_add<T>(input_tensor, index_tensor, update_tensor),
+                DataType::FLOAT32),
+        output_dim);
+  } else {
+    std::vector<int64_t> input_dim;
+    int x_dims = 1;
+    for (size_t i = 0; i < x.shape().size(); i++) {
+      x_dims *= x.shape()[i];
+    }
 
-  std::vector<int64_t> input_dim;
-  int x_dims = 1;
-  for (size_t i = 0; i < x.shape().size(); i++) {
-    x_dims *= x.shape()[i];
-  }
-
-  input_dim.push_back(x_dims);
-  input_dim.push_back(num_classes_tensor.shape()[0]);
-  auto input_tensor = full<T>(input_dim, 0, x.dtype());
+    input_dim.push_back(x_dims);
+    input_dim.push_back(num_classes_tensor.shape()[0]);
+    auto input_tensor = full<T>(input_dim, 0, x.dtype());
 
-  std::vector<int64_t> output_dim;
-  for (size_t i = 0; i < x.shape().size(); i++) {
-    output_dim.push_back(x.shape()[i]);
-  }
-  output_dim.push_back(num_classes_tensor.shape()[0]);
+    std::vector<int64_t> output_dim;
+    for (size_t i = 0; i < x.shape().size(); i++) {
+      output_dim.push_back(x.shape()[i]);
+    }
+    output_dim.push_back(num_classes_tensor.shape()[0]);
 
-  auto end = full<T>({1}, x_dims, x.dtype());
-  auto start = full<T>({1}, 0, x.dtype());
-  auto step = full<T>({1}, 1, x.dtype());
-  auto arange_tensor =
-      backend::arange_with_tensor<T>(start, end, step, x.dtype());
+    auto end = full<T>({1}, x_dims, x.dtype());
+    auto start = full<T>({1}, 0, x.dtype());
+    auto step = full<T>({1}, 1, x.dtype());
+    auto arange_tensor =
+        backend::arange_with_tensor<T>(start, end, step, x.dtype());
 
-  std::vector<int64_t> reshape_dim{x_dims, 1};
-  auto x_reshape = reshape<T>(x, reshape_dim);
-  auto arange_tensor_reshape = reshape<T>(arange_tensor, reshape_dim);
+    std::vector<int64_t> reshape_dim{x_dims, 1};
+    auto x_reshape = reshape<T>(x, reshape_dim);
+    auto arange_tensor_reshape = reshape<T>(arange_tensor, reshape_dim);
 
-  std::vector<Tensor> index_concat;
-  index_concat.push_back(arange_tensor_reshape);
-  index_concat.push_back(x_reshape);
-  auto index_tensor = concat<T>(index_concat, 1);
+    std::vector<Tensor> index_concat;
+    index_concat.push_back(arange_tensor_reshape);
+    index_concat.push_back(x_reshape);
+    auto index_tensor = concat<T>(index_concat, 1);
 
-  auto update_tensor = full<T>({x_dims}, 1, x.dtype());
+    auto update_tensor = full<T>({x_dims}, 1, x.dtype());
 
-  auto ans = reshape<T>(
-      cast<T>(scatter_nd_add<T>(input_tensor, index_tensor, update_tensor),
-              DataType::FLOAT32),
-      output_dim);
+    ans = reshape<T>(
+        cast<T>(scatter_nd_add<T>(input_tensor, index_tensor, update_tensor),
+                DataType::FLOAT32),
+        output_dim);
+  }
   return ans;
 }
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -300,6 +300,10 @@ def unstack_net6(x):
     return paddle.unstack(x, axis=-1)
 
 
+def one_hot_net(x):
+    return paddle.nn.functional.one_hot(x, 10)
+
+
 class TestPrimBase(unittest.TestCase):
     def setUp(self):
         np.random.seed(2023)
@@ -333,7 +337,7 @@ def base_net(self, flag=None):
             ops = [
                 op.name()
                 for op in fn.get_concrete_program(x)[-1]
-                .program.backward_program.global_block()
+                .program.forward_program.global_block()
                 .ops
             ]
             assert self.necessary_ops not in ops
@@ -504,7 +508,7 @@ def base_net(self, flag=None):
             ops = [
                 op.name()
                 for op in fn.get_concrete_program(x, y)[-1]
-                .program.backward_program.global_block()
+                .program.forward_program.global_block()
                 .ops
             ]
             assert self.necessary_ops not in ops
@@ -809,7 +813,7 @@ def base_net(self, flag=None):
             ops = [
                 op.name()
                 for op in fn.get_concrete_program(x, y, z)[-1]
-                .program.backward_program.global_block()
+                .program.forward_program.global_block()
                 .ops
             ]
             assert self.necessary_ops not in ops
@@ -1428,5 +1432,20 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimOneHot(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [50, 10, 70, 30]
+        self.dtype_x = "int32"
+        self.init_x_shape = [None, None, None, 30]
+        x = [np.random.randint(0, 9) for i in range(np.prod(self.shape_x))]
+        self.x = np.array(x).astype("int32").reshape(self.shape_x)
+        self.net = one_hot_net
+        self.necessary_ops = "pd_op.one_hot"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()
