[CINN]add Tril(u)Indices shape inference (#63000)

Xinyu302 · web-flow · commit f140f1ec0090 · 2024-03-27T20:11:15.000+08:00
* add Tril(u)Indices

* Update nullary_infer_sym.cc
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc
@@ -296,14 +296,64 @@ bool RandintOpInferSymbolicShape(
 
 bool TrilIndicesOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
+  const auto &attributes = op->attributes();
+  int rows = attributes.at("rows").dyn_cast<pir::Int32Attribute>().data();
+  int cols = attributes.at("cols").dyn_cast<pir::Int32Attribute>().data();
+  int offset = attributes.at("offset").dyn_cast<pir::Int32Attribute>().data();
+
+  const auto &out_sym_shape = [&] {
+    std::vector<symbol::DimExpr> out_sym_shape;
+    auto n_first_row =
+        offset > 0 ? std::min<int64_t>(cols, 1 + offset) : rows + offset > 0;
+    auto n_last_row =
+        std::max<int64_t>(0, std::min<int64_t>(cols, rows + offset));
+    auto n_row_all =
+        std::max<int64_t>(0, std::min<int64_t>(rows, rows + offset));
+    auto n_row_trapezoid = (n_last_row - n_first_row + 1);
+    auto tril_size = (n_first_row + n_last_row) * n_row_trapezoid >> 1;
+    auto diff_row = n_row_all - n_row_trapezoid;
+    if (diff_row > 0) {
+      tril_size += diff_row * cols;
+    }
+    out_sym_shape.emplace_back(std::int64_t(2));
+    out_sym_shape.emplace_back(std::int64_t(tril_size));
+    return out_sym_shape;
+  }();
+
+  symbol::ShapeOrDataDimExprs shape_data{
+      symbol::TensorShapeOrDataDimExprs(out_sym_shape)};
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   return true;
 }
 bool TriuIndicesOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
+  const auto &attributes = op->attributes();
+  int row = attributes.at("row").dyn_cast<pir::Int32Attribute>().data();
+  int col = attributes.at("col").dyn_cast<pir::Int32Attribute>().data();
+  int offset = attributes.at("offset").dyn_cast<pir::Int32Attribute>().data();
+
+  const auto &out_sym_shape = [&] {
+    std::vector<symbol::DimExpr> out_sym_shape;
+    offset = offset - 1;
+    auto n_first_row =
+        offset > 0 ? std::min<int64_t>(col, 1 + offset) : row + offset > 0;
+    auto n_last_row =
+        std::max<int64_t>(0, std::min<int64_t>(col, row + offset));
+    auto n_row_all = std::max<int64_t>(0, std::min<int64_t>(row, row + offset));
+    auto n_row_trapezoid = (n_last_row - n_first_row + 1);
+    auto tril_size = (n_first_row + n_last_row) * n_row_trapezoid >> 1;
+    auto diff_row = n_row_all - n_row_trapezoid;
+    if (diff_row > 0) {
+      tril_size += diff_row * col;
+    }
+    out_sym_shape.emplace_back(std::int64_t(2));
+    out_sym_shape.emplace_back(std::int64_t(row * col - tril_size));
+    return out_sym_shape;
+  }();
+
+  symbol::ShapeOrDataDimExprs shape_data{
+      symbol::TensorShapeOrDataDimExprs(out_sym_shape)};
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   return true;
 }
 bool UniformOpInferSymbolicShape(
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_nullary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_nullary_op.py
@@ -118,6 +118,92 @@ def test_eval_symbolic(self):
         return True
 
 
+class TriuIndicesNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        out = paddle.triu_indices(row=10, col=10, offset=0)
+        out = paddle.triu_indices(row=10, col=10, offset=2)
+        out = paddle.triu_indices(row=10, col=10, offset=-2)
+        out = paddle.triu_indices(row=10, col=3, offset=0)
+        out = paddle.triu_indices(row=10, col=3, offset=2)
+        out = paddle.triu_indices(row=10, col=3, offset=-2)
+        out = paddle.triu_indices(row=3, col=10, offset=0)
+        out = paddle.triu_indices(row=3, col=10, offset=2)
+        out = paddle.triu_indices(row=3, col=10, offset=-2)
+        return out
+
+
+class TriuIndicesOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.expected = [
+            'shape[2, 55], data[NULL]',
+            'shape[2, 36], data[NULL]',
+            'shape[2, 72], data[NULL]',
+            'shape[2, 6], data[NULL]',
+            'shape[2, 1], data[NULL]',
+            'shape[2, 12], data[NULL]',
+            'shape[2, 27], data[NULL]',
+            'shape[2, 21], data[NULL]',
+            'shape[2, 30], data[NULL]',
+        ]
+
+    def test_eval_symbolic(self):
+        net = TriuIndicesNet()
+        x_spec = InputSpec(shape=[None, None, None], dtype='float32')
+        input_spec = [x_spec]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+        check_infer_results(
+            net, input_spec, 'pd_op.triu_indices', self.expected
+        )
+        return True
+
+
+class TrilIndicesNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        out = paddle.tril_indices(row=10, col=10, offset=0)
+        out = paddle.tril_indices(row=10, col=10, offset=2)
+        out = paddle.tril_indices(row=10, col=10, offset=-2)
+        out = paddle.tril_indices(row=10, col=3, offset=0)
+        out = paddle.tril_indices(row=10, col=3, offset=2)
+        out = paddle.tril_indices(row=10, col=3, offset=-2)
+        out = paddle.tril_indices(row=3, col=10, offset=0)
+        out = paddle.tril_indices(row=3, col=10, offset=2)
+        out = paddle.tril_indices(row=3, col=10, offset=-2)
+        return out
+
+
+class TrilIndicesOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.expected = [
+            'shape[2, 55], data[NULL]',
+            'shape[2, 72], data[NULL]',
+            'shape[2, 36], data[NULL]',
+            'shape[2, 27], data[NULL]',
+            'shape[2, 30], data[NULL]',
+            'shape[2, 21], data[NULL]',
+            'shape[2, 6], data[NULL]',
+            'shape[2, 12], data[NULL]',
+            'shape[2, 1], data[NULL]',
+        ]
+
+    def test_eval_symbolic(self):
+        net = TrilIndicesNet()
+        x_spec = InputSpec(shape=[None, None, None], dtype='float32')
+        input_spec = [x_spec]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+        check_infer_results(
+            net, input_spec, 'pd_op.tril_indices', self.expected
+        )
+        return True
+
+
 class GaussianNet(paddle.nn.Layer):
     def __init__(self):
         super().__init__()
