cinn(op): fix slice symbolic shape (#62997)

6clc · web-flow · commit 564e10dcc090 · 2024-03-26T20:21:37.000+08:00
diff --git a/paddle/cinn/hlir/pe/transform.cc b/paddle/cinn/hlir/pe/transform.cc
@@ -1070,18 +1070,25 @@ ir::Tensor SliceSymbolic(const ir::Tensor& A,
     input_shape.emplace_back(shape);
   }
 
-  std::vector<int> new_starts(starts);
+  std::vector<Expr> new_starts;
+  std::transform(starts.begin(),
+                 starts.end(),
+                 std::back_inserter(new_starts),
+                 [](const int start) { return ir::Expr(start); });
+
   for (int i = 0; i < axes.size(); i++) {
-    CHECK(input_shape[axes[i]].is_constant())
-        << "Not supported Slice in dynamic dimensions, because the "
-           "relationship between slice range and symbol size cannot be "
-           "determined at compile time";
-    if (new_starts[i] < -input_shape[axes[i]].as_int64()) {
-      new_starts[i] = 0;
-    } else if (new_starts[i] < 0) {
-      new_starts[i] = input_shape[axes[i]].as_int64() + new_starts[i];
-    } else if (new_starts[i] > input_shape[axes[i]].as_int64()) {
-      new_starts[i] = input_shape[axes[i]].as_int64() - 1;
+    if (input_shape[axes[i]].is_constant()) {
+      if (new_starts[i].as_int64() < -input_shape[axes[i]].as_int64()) {
+        new_starts[i] = ir::Expr(0);
+      } else if (new_starts[i].as_int64() < 0) {
+        new_starts[i] = input_shape[axes[i]].as_int64() + new_starts[i];
+      } else if (new_starts[i].as_int64() > input_shape[axes[i]].as_int64()) {
+        new_starts[i] = input_shape[axes[i]].as_int64() - ir::Expr(1);
+      }
+    } else {
+      if (new_starts[i].as_int64() < 0) {
+        new_starts[i] = ir::Add::Make(input_shape[axes[i]], new_starts[i]);
+      }
     }
   }
 
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_multinary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_multinary_op.py
@@ -88,7 +88,7 @@ def test_eval_symbolic(self):
             )
 
             input_spec = [x_spec]
-            net = apply_to_static(net, True, input_spec)
+            net = apply_to_static(net, False, input_spec)
             net.eval()
             check_infer_results(net, input_spec, 'pd_op.slice', self.expected)
 

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ def test_eval_symbolic(self):`
`88`	`88`	`)`
`89`	`89`
`90`	`90`	`input_spec = [x_spec]`
`91`		`- net = apply_to_static(net, True, input_spec)`
	`91`	`+ net = apply_to_static(net, False, input_spec)`
`92`	`92`	`net.eval()`
`93`	`93`	`check_infer_results(net, input_spec, 'pd_op.slice', self.expected)`
`94`	`94`