[Mlir-commits] [mlir] [MLIR][NFC] Speed up is valid symbol check (PR #154924)

[William Moses]

https://github.com/wsmoses updated https://github.com/llvm/llvm-project/pull/154924

>From 19530291b90a28a9776eacd4369a8d371cb19708 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Fri, 22 Aug 2025 05:46:29 -0500
Subject: [PATCH] [MLIR][NFC] Speed up is valid symbol check

---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 47 +++++++++++++++---------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 22608a16cc1ab..8885286f8330a 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -427,6 +427,25 @@ bool mlir::affine::isValidSymbol(Value value) {
   return false;
 }
 
+/// A utility function to check if a value is defined at the top level of
+/// `region` or is an argument of `region` or dominates the region.
+static bool isTopLevelValueOrDominator(Value value, Region *region) {
+  Region *parentRegion;
+  if (auto arg = dyn_cast<BlockArgument>(value))
+    parentRegion = arg.getParentRegion();
+  else
+    parentRegion = value.getDefiningOp()->getParentRegion();
+  do {
+    if (parentRegion == region)
+      return true;
+    Operation *regionOp = region->getParentOp();
+    if (regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
+      break;
+    region = region->getParentOp()->getParentRegion();
+  } while (region);
+  return false;
+}
+
 /// A value can be used as a symbol for `region` iff it meets one of the
 /// following conditions:
 /// *) It is a constant.
@@ -445,17 +464,11 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
     return false;
 
   // A top-level value is a valid symbol.
-  if (region && ::isTopLevelValue(value, region))
+  if (region && isTopLevelValueOrDominator(value, region))
     return true;
 
   auto *defOp = value.getDefiningOp();
   if (!defOp) {
-    // A block argument that is not a top-level value is a valid symbol if it
-    // dominates region's parent op.
-    Operation *regionOp = region ? region->getParentOp() : nullptr;
-    if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
-      if (auto *parentOpRegion = region->getParentOp()->getParentRegion())
-        return isValidSymbol(value, parentOpRegion);
     return false;
   }
 
@@ -465,22 +478,22 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
     return true;
 
   // `Pure` operation that whose operands are valid symbolic identifiers.
-  if (isPure(defOp) && llvm::all_of(defOp->getOperands(), [&](Value operand) {
-        return affine::isValidSymbol(operand, region);
-      })) {
-    return true;
+  if (isPure(defOp)) {
+    bool allValid = true;
+    for (auto operand : defOp->getOperands()) {
+      if (!affine::isValidSymbol(operand, region)) {
+        allValid = false;
+        break;
+      }
+    }
+    if (allValid)
+      return true;
   }
 
   // Dim op results could be valid symbols at any level.
   if (auto dimOp = dyn_cast<ShapedDimOpInterface>(defOp))
     return isDimOpValidSymbol(dimOp, region);
 
-  // Check for values dominating `region`'s parent op.
-  Operation *regionOp = region ? region->getParentOp() : nullptr;
-  if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
-    if (auto *parentRegion = region->getParentOp()->getParentRegion())
-      return isValidSymbol(value, parentRegion);
-
   return false;
 }