[mlir][scf] Validate pipelining annotations in test pass (#208788)

I made the SCF pipelining test pass validate malformed scheduling
annotations before building the schedule. Bad annotations now produce
diagnostics instead of crashing, while valid schedules keep using the
existing path.
Fixes #206918

mlir/test/Dialect/SCF/loop-pipelining.mlir

@@ -1024,6 +1024,102 @@ func.func @invalid_schedule3(%A: memref<?xf32>, %result: memref<?xf32>, %ext: in
 
 // -----
 
+func.func @invalid_schedule_order_out_of_range(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    // expected-error@+1 {{invalid pipeline op order}}
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 3 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 1 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
+func.func @invalid_schedule_order_negative(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    // expected-error@+1 {{invalid pipeline op order}}
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = -1 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
+func.func @invalid_schedule_order_overflow(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    // expected-error@+1 {{invalid pipeline op order}}
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 9223372036854775808 : i128 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
+func.func @invalid_schedule_order_duplicate(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 0 } : memref<?xf32>
+    // expected-error@+1 {{duplicate pipeline op order}}
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
+func.func @invalid_schedule_stage_negative(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    // expected-error@+1 {{invalid pipeline stage}}
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = -1, __test_pipelining_op_order__ = 0 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
+func.func @invalid_schedule_stage_overflow(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c2 step %c1 {
+    // expected-error@+1 {{invalid pipeline stage}}
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 4294967296 : i128, __test_pipelining_op_order__ = 0 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+  } { __test_pipelining_loop__ }
+  return
+}
+
+// -----
+
 // Ensure this case not crash when step is zero.
 
 // CHECK-LABEL: @invalid_loop_step

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

@@ -21,6 +21,8 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+#include <limits>
+
 using namespace mlir;
 
 namespace {
@@ -154,18 +156,36 @@ struct TestSCFPipeliningPass
       return;
 
     schedule.resize(forOp.getBody()->getOperations().size() - 1);
-    forOp.walk([&schedule](Operation *op) {
+    WalkResult result = forOp.walk([&schedule](Operation *op) {
       auto attrStage =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningStageMarker);
       auto attrCycle =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningOpOrderMarker);
       if (attrCycle && attrStage) {
-        // TODO: Index can be out-of-bounds if ops of the loop body disappear
-        // due to folding.
-        schedule[attrCycle.getInt()] =
-            std::make_pair(op, unsigned(attrStage.getInt()));
+        const APInt &stage = attrStage.getValue();
+        if (stage.isNegative() ||
+            stage.getActiveBits() > std::numeric_limits<unsigned>::digits) {
+          op->emitOpError("invalid pipeline stage");
+          return WalkResult::interrupt();
+        }
+        const APInt &order = attrCycle.getValue();
+        if (order.isNegative() || order.getActiveBits() > 64 ||
+            order.getZExtValue() >= schedule.size()) {
+          op->emitOpError("invalid pipeline op order");
+          return WalkResult::interrupt();
+        }
+        size_t orderIndex = static_cast<size_t>(order.getZExtValue());
+        if (schedule[orderIndex].first) {
+          op->emitOpError("duplicate pipeline op order");
+          return WalkResult::interrupt();
+        }
+        schedule[orderIndex] =
+            std::make_pair(op, static_cast<unsigned>(stage.getZExtValue()));
       }
+      return WalkResult::advance();
     });
+    if (result.wasInterrupted())
+      schedule.clear();
   }
 
   /// Helper to generate "predicated" version of `op`. For simplicity we just