【PIR】add sequence_mask in pir (#59348)

yangguohao · kangguangli · web-flow · commit 669a3007e45b · 2023-12-08T17:16:03.000+08:00
* add sequence_mask in pir

* fix 123

* fix

* fix 2023-11-28

* fix 2023-11-29

* fix

* fix 2023-11-30

* fix

* fix3

* fix 1201

* fix 20231204

* fix 20231205

* fix

---------

Co-authored-by: kangguangli &lt;kangguangli@hotmail.com&gt;
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -99,6 +99,7 @@
     'self_dp_attention',
     'get_tensor_from_selected_rows',
     'print',
+    'sequence_mask',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
diff --git a/paddle/fluid/pir/dialect/operator/ir/update_ops.yaml b/paddle/fluid/pir/dialect/operator/ir/update_ops.yaml
@@ -15,3 +15,12 @@
     data_type : dtype
     backend : place
   support_tensor : [start, end, step]
+
+- op : sequence_mask
+  args: (Tensor x, Scalar(int) max_len, int out_dtype)
+  output: Tensor(y)
+  infer_meta:
+    func: SequenceMaskScalarInferMeta
+  kernel:
+    func: sequence_mask_scalar
+    data_type : x
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -2717,6 +2717,18 @@
 - op : send_uv (graph_send_uv)
   backward : send_uv_grad (graph_send_uv_grad)
 
+- op : sequence_mask
+  inputs:
+    x : X
+  attrs:
+    max_len: maxlen
+  outputs:
+    y : Y
+  scalar :
+    max_len :
+      data_type : int
+      tensor_name : MaxLenTensor
+
 - op : sequence_softmax
   backward : sequence_softmax_grad
   extra :
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3998,6 +3998,18 @@ void SplitWithNumInferMeta(const MetaTensor& x,
   }
 }
 
+void SequenceMaskScalarInferMeta(const MetaTensor& x,
+                                 const Scalar& max_len,
+                                 int out_dtype,
+                                 MetaTensor* y) {
+  auto dim = phi::vectorize<int>(x.dims());
+  int maxlen = max_len.to<int>();
+  dim.push_back(maxlen > 0 ? maxlen : -1);
+  y->set_dims(phi::make_ddim(dim));
+  auto out_phi_dtype = phi::TransToPhiDataType(out_dtype);
+  y->set_dtype(out_phi_dtype);
+}
+
 void SquaredL2NormInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_dims({1});
 }
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -550,6 +550,11 @@ void RReluGradInferMeta(const MetaTensor& out_grad,
                         const MetaTensor& noise,
                         MetaTensor* x_grad);
 
+void SequenceMaskScalarInferMeta(const MetaTensor& x,
+                                 const Scalar& max_len,
+                                 int out_dtype,
+                                 MetaTensor* y);
+
 void SetValueInferMeta(const MetaTensor& x, MetaTensor* out);
 
 void ShapeInferMeta(const MetaTensor& input, MetaTensor* out);
diff --git a/paddle/phi/kernels/cpu/sequence_mask_kernel.cc b/paddle/phi/kernels/cpu/sequence_mask_kernel.cc
@@ -27,3 +27,13 @@ PD_REGISTER_KERNEL(sequence_mask,
                    int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
+PD_REGISTER_KERNEL(sequence_mask_scalar,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SequenceMaskScalarKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
diff --git a/paddle/phi/kernels/gpu/sequence_mask_kernel.cu b/paddle/phi/kernels/gpu/sequence_mask_kernel.cu
@@ -27,3 +27,13 @@ PD_REGISTER_KERNEL(sequence_mask,
                    int64_t) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
+PD_REGISTER_KERNEL(sequence_mask_scalar,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SequenceMaskScalarKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}
diff --git a/paddle/phi/kernels/impl/sequence_mask_kernel_impl.h b/paddle/phi/kernels/impl/sequence_mask_kernel_impl.h
@@ -27,6 +27,42 @@ limitations under the License. */
 
 namespace phi {
 
+template <typename T, typename Context>
+void SequenceMaskScalarKernel(const Context& ctx,
+                              const DenseTensor& x,
+                              const Scalar& max_len,
+                              int out_dtype,
+                              DenseTensor* y) {
+  int maxlen = max_len.to<int>();
+  auto* x_data = x.data<T>();
+  auto x_numel = x.numel();
+
+  if (maxlen < 0) {
+    if (x_numel == 0) {
+      maxlen = 0;
+    } else {
+#if defined(__NVCC__) || defined(__HIPCC__)
+      VLOG(10)
+          << "SequenceMaskOp on GPU may be slow when maxlen is not provided.";
+      maxlen = static_cast<int>(
+          thrust::reduce(thrust::device_pointer_cast(x_data),
+                         thrust::device_pointer_cast(x_data) + x_numel,
+                         static_cast<T>(0),
+                         thrust::maximum<T>()));
+#else
+      maxlen = static_cast<int>(*std::max_element(x_data, x_data + x_numel));
+#endif
+    }
+    auto y_dim = common::vectorize<int>(x.dims());
+    y_dim.push_back(maxlen);
+    y->Resize(common::make_ddim(y_dim));
+  }
+
+  phi::VisitDataType(phi::TransToPhiDataType(out_dtype),
+                     phi::funcs::SequenceMaskFunctor<Context, T>(
+                         ctx, x_data, y, x_numel * maxlen, maxlen));
+}
+
 template <typename T, typename Context>
 void SequenceMaskKernel(const Context& ctx,
                         const DenseTensor& x,
@@ -56,33 +92,6 @@ void SequenceMaskKernel(const Context& ctx,
             "received Input(MaxLenTensor) value = %d.",
             maxlen));
   }
-
-  auto* x_data = x.data<T>();
-  auto x_numel = x.numel();
-
-  if (maxlen < 0) {
-    if (x_numel == 0) {
-      maxlen = 0;
-    } else {
-#if defined(__NVCC__) || defined(__HIPCC__)
-      VLOG(10)
-          << "SequenceMaskOp on GPU may be slow when maxlen is not provided.";
-      maxlen = static_cast<int>(
-          thrust::reduce(thrust::device_pointer_cast(x_data),
-                         thrust::device_pointer_cast(x_data) + x_numel,
-                         static_cast<T>(0),
-                         thrust::maximum<T>()));
-#else
-      maxlen = static_cast<int>(*std::max_element(x_data, x_data + x_numel));
-#endif
-    }
-    auto y_dim = common::vectorize<int>(x.dims());
-    y_dim.push_back(maxlen);
-    y->Resize(common::make_ddim(y_dim));
-  }
-
-  phi::VisitDataType(phi::TransToPhiDataType(out_dtype),
-                     phi::funcs::SequenceMaskFunctor<Context, T>(
-                         ctx, x_data, y, x_numel * maxlen, maxlen));
+  SequenceMaskScalarKernel<T, Context>(ctx, x, Scalar(maxlen), out_dtype, y);
 }
 }  // namespace phi
diff --git a/test/dygraph_to_static/dygraph_to_static_utils.py b/test/dygraph_to_static/dygraph_to_static_utils.py
@@ -100,11 +100,8 @@ def lower_case_name(self):
 DISABLED_IR_TEST_FILES = {
     IrMode.LEGACY_IR: [],
     IrMode.PT: [
-        "test_gradname_parse",
-        "test_seq2seq",
         "test_save_inference_model",
         "test_tensor_hook",
-        "test_reinforcement_learning",
     ],
     IrMode.PIR: [],
 }
diff --git a/test/sequence/CMakeLists.txt b/test/sequence/CMakeLists.txt
@@ -10,3 +10,12 @@ endforeach()
 set_tests_properties(test_sequence_conv PROPERTIES TIMEOUT 120)
 set_tests_properties(test_sequence_concat PROPERTIES TIMEOUT 120)
 set_tests_properties(test_sequence_pool PROPERTIES TIMEOUT 120)
+
+set(PIR_COVERAGE_TESTS test_sequence_mask)
+
+foreach(PIR_COVERAGE_TEST ${PIR_COVERAGE_TESTS})
+  py_test_modules(${PIR_COVERAGE_TEST}_pir MODULES ${PIR_COVERAGE_TEST} ENVS
+                  FLAGS_enable_pir_in_executor=true)
+  set_tests_properties(${PIR_COVERAGE_TEST}_pir PROPERTIES TIMEOUT 120)
+  message(STATUS "PIR Copied OpTest: ${PIR_COVERAGE_TEST}_pir in sequence test")
+endforeach()

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,7 @@`
`99`	`99`	`'self_dp_attention',`
`100`	`100`	`'get_tensor_from_selected_rows',`
`101`	`101`	`'print',`
	`102`	`+ 'sequence_mask',`
`102`	`103`	`]`
`103`	`104`
`104`	`105`	`NO_NEED_GEN_STATIC_ONLY_APIS = [`
Original file line number	Diff line number	Diff line change
`@@ -3998,6 +3998,18 @@ void SplitWithNumInferMeta(const MetaTensor& x,`
`3998`	`3998`	`}`
`3999`	`3999`	`}`
`4000`	`4000`
	`4001`	`+void SequenceMaskScalarInferMeta(const MetaTensor& x,`
	`4002`	`+ const Scalar& max_len,`
	`4003`	`+ int out_dtype,`
	`4004`	`+ MetaTensor* y) {`
	`4005`	`+ auto dim = phi::vectorize<int>(x.dims());`
	`4006`	`+ int maxlen = max_len.to<int>();`
	`4007`	`+ dim.push_back(maxlen > 0 ? maxlen : -1);`
	`4008`	`+ y->set_dims(phi::make_ddim(dim));`
	`4009`	`+ auto out_phi_dtype = phi::TransToPhiDataType(out_dtype);`
	`4010`	`+ y->set_dtype(out_phi_dtype);`
	`4011`	`+}`
	`4012`	`+`
`4001`	`4013`	`void SquaredL2NormInferMeta(const MetaTensor& x, MetaTensor* out) {`
`4002`	`4014`	`out->set_dims({1});`
`4003`	`4015`	`}`