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

[Oleksandr Alex Zinenko]

================
@@ -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)
+    return failure();
+
+  IRRewriter b(forOp->getContext());
----------------
ftynse wrote:

It is generally preferrable to take a `RewriterBase &` as argument than to create a local one.

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