【PIR OpTest Fix No.6】 fix test_tdm_sampler_op

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -131,6 +131,7 @@
     'send_v2',
     'shadow_feed',
     'sparse_momentum',
+    'tdm_sampler',
     'soft_relu',
     'uniform_random_batch_size_like',
 ]

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -1243,6 +1243,16 @@
   inplace : (mean -> mean_out), (variance -> variance_out)
   optional : reserve_space
 
+- op : tdm_sampler
+  args: (Tensor x, Tensor travel, Tensor layer, bool output_positive=true, int[] neg_samples_num_list={}, int[] layer_offset_lod={}, int seed = 0, int dtype=2)
+  output: Tensor(out), Tensor(labels), Tensor(mask)
+  infer_meta:
+    func : TdmSamplerInferMeta
+  kernel:
+    func : tdm_sampler
+    data_type : x
+  optional : labels
+
 - op : tile
   args : (Tensor x, IntArray repeat_times = {})
   output : Tensor(out)

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -52,6 +52,7 @@ const std::unordered_set<std::string> LegacyOpList = {
     ShareDataOp::name(),
     SparseMomentumOp::name(),
     GetTensorFromSelectedRowsOp::name(),
+    TdmSamplerOp::name(),
     RowConvOp::name(),
     RowConvGradOp::name(),
     SoftReluOp::name(),

paddle/phi/api/yaml/op_compat.yaml

@@ -3109,6 +3109,12 @@
   extra :
     attrs : [bool use_mkldnn = false, bool use_cudnn = false]
 
+- op : tdm_sampler
+  inputs:
+    {x : X, travel : Travel, layer : Layer}
+  outputs:
+    {out : Out, labels : Labels, mask : Mask}
+
 - op : thresholded_relu
   inputs :
     x : X

paddle/phi/infermeta/ternary.cc

@@ -1525,5 +1525,46 @@ void QuantLinearInferMeta(const MetaTensor& x,
   y->share_lod(x);
   y->set_dtype(x.dtype());
 }
+void TdmSamplerInferMeta(const MetaTensor& x,
+                         const MetaTensor& travel,
+                         const MetaTensor& layer,
+                         bool output_positive,
+                         const std::vector<int>& neg_samples_num_list,
+                         const std::vector<int>& layer_offset_lod,
+                         int seed,
+                         int dtype,
+                         MetaTensor* out,
+                         MetaTensor* labels,
+                         MetaTensor* mask,
+                         MetaConfig config) {
+  auto neg_samples_num_vec = neg_samples_num_list;
+  auto output_positive_flag = output_positive;
 
+  int64_t sample_res_length = 0;
+  for (auto sample_nums : neg_samples_num_vec) {
+    sample_res_length += sample_nums + (int64_t)output_positive_flag;
+  }
+  auto ddim = phi::make_ddim({-1, sample_res_length});
+  auto input_dims = x.dims();
+  if (config.is_runtime) {
+    auto output_dims = phi::vectorize(input_dims);
+    auto batch_size = output_dims[0];
+    out->set_dims(phi::make_ddim({batch_size, sample_res_length}));
+    mask->set_dims(phi::make_ddim({batch_size, sample_res_length}));
+    if (labels) {
+      labels->set_dims(phi::make_ddim({batch_size, sample_res_length}));
+    }
+  } else {
+    out->set_dims(ddim);
+    mask->set_dims(ddim);
+    if (labels) {
+      labels->set_dims(ddim);
+    }
+  }
+  out->set_dtype(x.dtype());
+  mask->set_dtype(x.dtype());
+  if (labels) {
+    labels->set_dtype(x.dtype());
+  }
+}
 }  // namespace phi

paddle/phi/infermeta/ternary.h

@@ -238,4 +238,17 @@ void QuantLinearInferMeta(const MetaTensor& x,
                           float quant_min_bound,
                           MetaTensor* y);
 
+void TdmSamplerInferMeta(const MetaTensor& x,
+                         const MetaTensor& travel,
+                         const MetaTensor& layer,
+                         bool output_positive,
+                         const std::vector<int>& neg_samples_num_list,
+                         const std::vector<int>& layer_offset_lod,
+                         int seed,
+                         int dtype,
+                         MetaTensor* out,
+                         MetaTensor* labels,
+                         MetaTensor* mask,
+                         MetaConfig config = MetaConfig());
+
 }  // namespace phi

test/white_list/pir_op_test_no_check_list

@@ -7,5 +7,6 @@ test_poisson_op
 test_randint_op
 test_randperm_op
 test_seed_op
+test_tdm_sampler_op
 test_uniform_random_bf16_op
 test_uniform_random_inplace_op

test/white_list/pir_op_test_white_list

@@ -290,6 +290,7 @@ test_sum_bf16_mkldnn_op
 test_sum_mkldnn_op
 test_svd_op
 test_take_along_axis_op
+test_tdm_sampler_op
 test_temporal_shift_op
 test_tile_op
 test_top_k_v2_op