[CINN] Add InferSymbolicShapeInterface for (frame, prune_gate_by_capacity) op (#68644)

ooooo-create · web-flow · commit 82bea0338b6f · 2024-11-06T15:31:48.000+08:00
* add infersymbolicshape interface for frame &amp;&amp; prune_gate_by_capacity

* fix compile

* Fix compile

* fix ci

* Fix

* Apply review

* Apply review
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -1725,6 +1725,28 @@ bool PriorBoxOpInferSymbolicShape(
   return true;
 }
 
+bool PruneGateByCapacityOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &expert_count_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &expert_count_shape = expert_count_shape_or_data.shape();
+  int64_t n_expert = op->attribute<pir::Int64Attribute>("n_expert").data();
+  int64_t n_worker = op->attribute<pir::Int64Attribute>("n_worker").data();
+  symbol::DimExpr expert_count_num_ele = 1;
+  for (const auto &i : expert_count_shape) {
+    expert_count_num_ele = expert_count_num_ele * i;
+  }
+
+  infer_context->AddEqualCstr(expert_count_num_ele,
+                              symbol::DimExpr(n_expert * n_worker));
+  const auto &gate_idx_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(gate_idx_shape_or_data.shape())});
+  return true;
+}
 // bool PullBoxSparseOpInferSymbolicShape(pir::Operation *op,
 //                                        pir::InferSymbolicShapeContext
 //                                        *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -79,6 +79,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(MarginCrossEntropy)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatmulWithFlatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(PriorBox)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(PruneGateByCapacity)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullGpuPsSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullSparseV2)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -1087,6 +1087,77 @@ bool FakeChannelWiseQuantizeDequantizeAbsMaxOpInferSymbolicShape(
   return true;
 }
 
+bool FrameOpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &x_shape = x_shape_or_data.shape();
+  size_t x_rank = x_shape.size();
+  PADDLE_ENFORCE_GE(x_rank,
+                    1,
+                    common::errors::InvalidArgument(
+                        "Input(X) of FrameOp should be a tensor which contains "
+                        "at least 1 dimension, but got rank %s.",
+                        x_rank));
+  int hop_length = op->attribute<pir::Int32Attribute>("hop_length").data();
+  int frame_length = op->attribute<pir::Int32Attribute>("frame_length").data();
+  int axis = op->attribute<pir::Int32Attribute>("axis").data();
+  PADDLE_ENFORCE_GT(hop_length,
+                    0,
+                    common::errors::InvalidArgument(
+                        "Attribute(hop_length) of FrameOp should be greater "
+                        "than 0, but got %s.",
+                        hop_length));
+  PADDLE_ENFORCE_EQ(
+      (axis == 0 || axis == -1),
+      true,
+      common::errors::InvalidArgument(
+          "Attribute(axis) of FrameOp should 0 or -1, but got %s.", axis));
+  std::vector<symbol::DimExpr> out_shape;
+  symbol::DimExpr seq_length;
+  symbol::DimExpr n_frames = 0;
+  int start_axis = 0;
+  int end_axis = 0;
+
+  if (axis == 0) {
+    seq_length = x_shape[0];
+    start_axis = 1;
+    end_axis = x_rank - 1;
+  } else {
+    seq_length = x_shape[x_rank - 1];
+    start_axis = 0;
+    end_axis = x_rank - 2;
+  }
+  if (seq_length.isa<int64_t>()) {
+    PADDLE_ENFORCE_LE(frame_length,
+                      seq_length.dyn_cast<int64_t>(),
+                      common::errors::InvalidArgument(
+                          "Attribute(frame_length) of FrameOp should be less "
+                          "equal than sequence length, but got (%s) > (%s).",
+                          frame_length,
+                          seq_length.dyn_cast<int64_t>()));
+  }
+
+  // It won't go into for loop when x_rank == 1U.
+  for (int i = start_axis; i <= end_axis; i++) {
+    out_shape.push_back(x_shape[i]);
+  }
+
+  n_frames = symbol::DimExpr((seq_length - frame_length) / hop_length + 1);
+  if (axis == 0) {
+    out_shape.insert(out_shape.begin(), symbol::DimExpr(frame_length));
+    out_shape.insert(out_shape.begin(), n_frames);
+  } else {
+    out_shape.push_back(symbol::DimExpr(frame_length));
+    out_shape.push_back(n_frames);
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
+
 bool EigOpInferSymbolicShape(pir::Operation *op,
                              pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -68,6 +68,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseQuantizeDequantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftR2c)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Frame)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2084,6 +2084,7 @@
   kernel :
     func : frame
   backward : frame_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : frobenius_norm
   args : (Tensor x, IntArray axis,  bool keep_dim,  bool reduce_all)
@@ -3856,6 +3857,7 @@
   kernel :
     func : prune_gate_by_capacity
     data_type : gate_idx
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
 
 - op : psroi_pool
diff --git a/test/legacy_test/test_frame_op.py b/test/legacy_test/test_frame_op.py
@@ -74,7 +74,7 @@ def init_attrs(self):
 
     def test_check_output(self):
         paddle.enable_static()
-        self.check_output()
+        self.check_output(check_pir=True)
         paddle.disable_static()
 
     def test_check_grad_normal(self):
diff --git a/test/legacy_test/test_prune_gate_by_capacity_op.py b/test/legacy_test/test_prune_gate_by_capacity_op.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from op_test import OpTest
 
 import paddle
 from paddle.base import core
@@ -65,6 +66,40 @@ def assert_allclose(output, expected, n_expert):
     np.testing.assert_allclose(c1, c2)
 
 
+@unittest.skipIf(
+    not core.is_compiled_with_cuda(), "core is not compiled with CUDA"
+)
+class TestPruneGateByCapacityOp(OpTest):
+    def _get_places(self):
+        return [paddle.CUDAPlace(0)]
+
+    def setUp(self):
+        self.op_type = "prune_gate_by_capacity"
+        self.python_api = utils._prune_gate_by_capacity
+        self.dtype = np.int64
+        self.n_expert = 24
+        self.n_worker = 2
+        self.gate_idx = np.random.randint(0, self.n_expert, size=(200,)).astype(
+            self.dtype
+        )
+        self.expert_count = count(self.gate_idx, self.n_expert * self.n_worker)
+        self.out = prune_gate_by_capacity(
+            self.gate_idx, self.expert_count, self.n_expert, self.n_worker
+        ).astype(self.dtype)
+        self.inputs = {
+            "GateIdx": self.gate_idx,
+            "ExpertCount": self.expert_count,
+        }
+        self.attrs = {
+            'n_expert': self.n_expert,
+            'n_worker': self.n_worker,
+        }
+        self.outputs = {"NewGateIdx": self.out}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda(), "core is not compiled with CUDA"
 )
