[PIR] Register some operators to pir (#62384)

zhangbo9674 · web-flow · commit a7e2a2db9ed0 · 2024-03-05T19:52:11.000+08:00
* add ops

* fix
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -115,6 +115,7 @@
     'quantize_linear_',
     'dequantize_linear',
     'dequantize_linear_',
+    'coalesce_tensor_',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
@@ -172,6 +173,8 @@
     'push_sparse_v2',
     'push_sparse_v2_',
     'partial_send',
+    'nop',
+    'nop_',
 ]
 
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc b/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc
@@ -32,6 +32,14 @@ KernelKeyTuple SaveCombineOpParseKernelKey(pir::Operation* op) {
   return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};
 }
 
+KernelKeyTuple NopOpParseKernelKey(pir::Operation* op) {
+  return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};
+}
+
+KernelKeyTuple Nop_OpParseKernelKey(pir::Operation* op) {
+  return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};
+}
+
 }  // namespace paddle::dialect
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ParseKernelKeyInterface)
diff --git a/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h b/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h
@@ -59,6 +59,10 @@ KernelKeyTuple UniqueOpParseKernelKey(pir::Operation *op);
 
 KernelKeyTuple SaveCombineOpParseKernelKey(pir::Operation *op);
 
+KernelKeyTuple NopOpParseKernelKey(pir::Operation *op);
+
+KernelKeyTuple Nop_OpParseKernelKey(pir::Operation *op);
+
 }  // namespace dialect
 }  // namespace paddle
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -331,6 +331,16 @@
     func : channel_shuffle
   backward : channel_shuffle_grad
 
+- op : coalesce_tensor_
+  args : (Tensor[] input, DataType dtype, bool copy_data = false, bool set_constant = false, bool persist_output = false, float constant = 0.0, bool use_align = true, int align_size = -1, int size_of_dtype = -1, int64_t[] concated_shapes = {}, int64_t[] concated_ranks = {})
+  output : Tensor[](output){input.size()}, Tensor(fused_output)
+  infer_meta :
+    func : CoalesceTensorInferMeta
+  kernel :
+    func : coalesce_tensor
+    data_type : dtype
+  inplace: (input -> output)
+
 - op : conv2d_transpose
   args : (Tensor x, Tensor filter, int[] strides={1, 1}, int[] paddings={0, 0}, int[] output_padding={}, IntArray output_size={}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1}, str data_format="NCHW")
   output : Tensor(out)
@@ -1049,6 +1059,16 @@
   backward : multiply_grad
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
+- op : nop
+  args : (Tensor x)
+  output : Tensor(out)
+  infer_meta :
+    func : UnchangedInferMeta
+  kernel :
+    func : nop
+  inplace: (x -> out)
+  interfaces : paddle::dialect::ParseKernelKeyInterface
+
 - op : norm
   args : (Tensor x, int axis, float epsilon, bool is_test)
   output : Tensor(out), Tensor(norm)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3767,6 +3767,12 @@
   outputs:
     {cost : Cost, sample_logits : SampleLogits, sample_labels : SampleLabels}
 
+- op: nop
+  inputs :
+    x : X
+  outputs :
+    out : Out
+
 - op: number_count
   inputs :
     {numbers: numbers}

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,7 @@`
`115`	`115`	`'quantize_linear_',`
`116`	`116`	`'dequantize_linear',`
`117`	`117`	`'dequantize_linear_',`
	`118`	`+ 'coalesce_tensor_',`
`118`	`119`	`]`
`119`	`120`
`120`	`121`	`NO_NEED_GEN_STATIC_ONLY_APIS = [`
`@@ -172,6 +173,8 @@`
`172`	`173`	`'push_sparse_v2',`
`173`	`174`	`'push_sparse_v2_',`
`174`	`175`	`'partial_send',`
	`176`	`+ 'nop',`
	`177`	`+ 'nop_',`
`175`	`178`	`]`
`176`	`179`
`177`	`180`