【Infer Symbolic Shape No.170】Add infer_symbol_shape for random_routing (#68670)

ooooo-create · web-flow · commit 13e7020d8744 · 2024-10-30T10:59:58.000+08:00
* test

* test

* add random_routing

* Fix

* Fix

* Fix

* fix
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -3267,6 +3267,28 @@ bool PyramidHashOpInferSymbolicShape(
 //   // pass
 //   return true;
 // }
+bool RandomRoutingOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &top_value_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+  const auto &prob_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &top_idx_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2)).shape();
+  infer_context->AddEqualCstr(prob_shape[0], top_value_shape[0]);
+  infer_context->AddEqualCstr(top_idx_shape[0], top_value_shape[0]);
+  infer_context->AddEqualCstr(top_idx_shape[1], top_value_shape[1]);
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(top_idx_shape)});
+  return true;
+}
+
+bool RandomRouting_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return RandomRoutingOpInferSymbolicShape(op, infer_context);
+}
 
 bool RmsNormOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -106,6 +106,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(PyramidHash)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(QuantizeLinear)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(QuantizeLinear_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(RankAttention)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(RandomRouting)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(RandomRouting_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RmsNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiPool)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Rnn)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3963,6 +3963,7 @@
     data_type : prob
   inplace : (topk_idx -> out)
   traits : pir::SideEffectTrait, paddle::dialect::ForwardOnlyTrait
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : randperm
   args : (int n, DataType dtype, Place place={})
