[Mlir-commits] [mlir] [mlir][affine] Add affine-loop-constantize-bounds transfrom pass to affine dialect (PR #211989)

[lonely eagle]

================
@@ -0,0 +1,104 @@
+//===- AffineLoopConstantizeBounds.cpp - Constantize loop bounds pass ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to transform affine loops with symbolic bounds
+// into a static-trip-count main loop and a residual tail loop using range
+// analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include <cstdint>
+
+namespace mlir {
+namespace affine {
+#define GEN_PASS_DEF_AFFINELOOPCONSTANTIZEBOUNDS
+#include "mlir/Dialect/Affine/Transforms/Passes.h.inc"
+} // namespace affine
+} // namespace mlir
+
+#define DEBUG_TYPE "affine-loop-constantize-bounds"
+
+using namespace mlir;
+using namespace mlir::affine;
+
+namespace {
+
+/// Computes the constant upper or lower bound for a given affine map expression
+/// \p map and its operands \p operands, constrained by the specified \p type.
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+                                               ValueRange operands,
+                                               presburger::BoundType type) {
+  ValueBoundsConstraintSet::Variable var(map, operands);
+  return ValueBoundsConstraintSet::computeConstantBound(type, var, nullptr,
+                                                        {true, true});
+}
+
+static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
+  // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
+  if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
+    return failure();
+  if (forOp.getStepAsInt() != 1)
+    return failure();
+
+  // Infer the range [upperMin, upperMax] for the upper bound. We require a
+  // strictly positive minimum bound (upperMin > 0) to guarantee a safe,
+  // non-empty static trip count for the main loop.
+  FailureOr<int64_t> upperMin = computeConstantBound(
+      forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+      presburger::BoundType::LB);
+  FailureOr<int64_t> upperMax = computeConstantBound(
+      forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+      presburger::BoundType::UB);
+  if (failed(upperMin) || *upperMin <= 0)
----------------
linuxlonelyeagle wrote:

If we set this to an assert, I'm afraid the case below will trigger it, even though it should currently be a valid IR. Considering the scenario below, I'm making it a dead condition.
```
  func.func @fully_constantized_no_peeling() {
    %c0 = arith.constant 0 : index
    %0 = test.value_with_bounds {max = 1 : index, min = 0 : index}
    %1 = affine.for %arg0 = 0 to -3 iter_args(%arg1 = %c0) -> (index) {
      %2 = affine.apply #map(%arg0)
      %3 = arith.addi %arg1, %0 : index
      affine.yield %3 : index
    }
    return
  }
```

https://github.com/llvm/llvm-project/pull/211989