[polly] [polly] Bound ISL operations during pre-vectorization (PR #143876)

[Shikhar Jain via llvm-commits]

https://github.com/ShikharjQUIC updated https://github.com/llvm/llvm-project/pull/143876

>From b6fce5961cc1860465ab1201d82d0f90c18ace94 Mon Sep 17 00:00:00 2001
From: ShikharJain <shikharj at qti.qualcomm.com>
Date: Fri, 23 May 2025 07:29:01 -0700
Subject: [PATCH] [polly] Bound ISL operations during pre-vectorization

Bound ISL operations during pre-vectorization to
prevent indefinite compilation. If the bound fails,
mark the SCoP invalid. Includes relevant unit test.

Fixes #132717.
---
 polly/lib/Transform/ScheduleOptimizer.cpp     | 23 ++++++++++-
 .../prevectorization_islbound.ll              | 41 +++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 polly/test/ScheduleOptimizer/prevectorization_islbound.ll

diff --git a/polly/lib/Transform/ScheduleOptimizer.cpp b/polly/lib/Transform/ScheduleOptimizer.cpp
index 070700a64a168..8bd4f20636cdc 100644
--- a/polly/lib/Transform/ScheduleOptimizer.cpp
+++ b/polly/lib/Transform/ScheduleOptimizer.cpp
@@ -387,6 +387,8 @@ ScheduleTreeOptimizer::isolateFullPartialTiles(isl::schedule_node Node,
   isl::union_set ScheduleRangeUSet = SchedRelUMap.range();
   isl::set ScheduleRange{ScheduleRangeUSet};
   isl::set IsolateDomain = getPartialTilePrefixes(ScheduleRange, VectorWidth);
+  if (IsolateDomain.is_null())
+    return isl::schedule_node();
   auto AtomicOption = getDimOptions(IsolateDomain.ctx(), "atomic");
   isl::union_set IsolateOption = getIsolateOptions(IsolateDomain, 1);
   Node = Node.parent().parent();
@@ -430,6 +432,8 @@ isl::schedule_node ScheduleTreeOptimizer::prevectSchedBand(
   Node =
       isl::manage(isl_schedule_node_band_tile(Node.release(), Sizes.release()));
   Node = isolateFullPartialTiles(Node, VectorWidth);
+  if (Node.is_null())
+    return isl::schedule_node();
   Node = Node.child(0);
   // Make sure the "trivially vectorizable loop" is not unrolled. Otherwise,
   // we will have troubles to match it in the backend.
@@ -543,6 +547,8 @@ ScheduleTreeOptimizer::applyPrevectBandOpt(isl::schedule_node Node) {
       break;
     }
 
+  if (Node.is_null())
+    return isl::schedule_node();
   return Node;
 }
 
@@ -574,7 +580,12 @@ ScheduleTreeOptimizer::optimizeBand(__isl_take isl_schedule_node *NodeArg,
   if (OAI->Prevect) {
     // FIXME: Prevectorization requirements are different from those checked by
     // isTileableBandNode.
-    Node = applyPrevectBandOpt(Node);
+    {
+      IslMaxOperationsGuard MaxOpGuard(Node.ctx().get(), ScheduleComputeOut);
+      Node = applyPrevectBandOpt(Node);
+      if (MaxOpGuard.hasQuotaExceeded() || Node.is_null())
+        return (isl::schedule_node()).release();
+    }
   }
 
   return Node.release();
@@ -585,6 +596,8 @@ ScheduleTreeOptimizer::optimizeSchedule(isl::schedule Schedule,
                                         const OptimizerAdditionalInfoTy *OAI) {
   auto Root = Schedule.get_root();
   Root = optimizeScheduleNode(Root, OAI);
+  if (Root.is_null())
+    return isl::schedule();
   return Root.get_schedule();
 }
 
@@ -593,6 +606,8 @@ isl::schedule_node ScheduleTreeOptimizer::optimizeScheduleNode(
   Node = isl::manage(isl_schedule_node_map_descendant_bottom_up(
       Node.release(), optimizeBand,
       const_cast<void *>(static_cast<const void *>(OAI))));
+  if (Node.is_null())
+    return isl::schedule_node();
   return Node;
 }
 
@@ -908,6 +923,12 @@ static void runIslScheduleOptimizer(
       DepsChanged};
   if (OAI.PatternOpts || OAI.Postopts || OAI.Prevect) {
     Schedule = ScheduleTreeOptimizer::optimizeSchedule(Schedule, &OAI);
+    if (Schedule.is_null()) {
+      POLLY_DEBUG(dbgs() << "Invalidating scop from further optimization as "
+                            "max operations exceeded\n");
+      S.invalidate(COMPLEXITY, DebugLoc());
+      return;
+    }
     Schedule = hoistExtensionNodes(Schedule);
     POLLY_DEBUG(printSchedule(dbgs(), Schedule, "After post-optimizations"));
     walkScheduleTreeForStatistics(Schedule, 2);
diff --git a/polly/test/ScheduleOptimizer/prevectorization_islbound.ll b/polly/test/ScheduleOptimizer/prevectorization_islbound.ll
new file mode 100644
index 0000000000000..bf6a763a616ef
--- /dev/null
+++ b/polly/test/ScheduleOptimizer/prevectorization_islbound.ll
@@ -0,0 +1,41 @@
+; RUN: opt %loadNPMPolly -S "-polly-vectorizer=stripmine" -polly-opt-isl -polly-debug -disable-output < %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+; ModuleID = '<stdin>'
+source_filename = "<stdin>"
+target datalayout = "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-i128:128-f64:32:64-f80:32-n8:16:32-S128"
+target triple = "i386-unknown-linux-gnu"
+
+define void @ham(ptr %arg, ptr %arg1, i32 %arg2, i32 %arg3, ptr %arg4, i32 %arg5, i32 %arg6) {
+bb:
+  %getelementptr = getelementptr [7 x float], ptr null, i32 0, i32 %arg3
+  br label %bb7
+
+bb7:                                              ; preds = %bb11, %bb
+  %phi = phi i32 [ 0, %bb ], [ %add16, %bb11 ]
+  br label %bb8
+
+bb8:                                              ; preds = %bb8, %bb7
+  %phi9 = phi i32 [ 0, %bb7 ], [ %add, %bb8 ]
+  %getelementptr10 = getelementptr [7 x float], ptr null, i32 0, i32 %phi9
+  store float 0.000000e+00, ptr %getelementptr10, align 4
+  %add = add i32 %phi9, 1
+  %icmp = icmp eq i32 %phi9, 0
+  br i1 %icmp, label %bb8, label %bb11
+
+bb11:                                             ; preds = %bb8
+  %load = load float, ptr %getelementptr, align 4
+  store float %load, ptr %arg4, align 4
+  %getelementptr12 = getelementptr [7 x float], ptr null, i32 0, i32 %arg5
+  %load13 = load float, ptr %getelementptr12, align 4
+  store float %load13, ptr %arg, align 4
+  %getelementptr14 = getelementptr [7 x float], ptr null, i32 0, i32 %arg6
+  %load15 = load float, ptr %getelementptr14, align 4
+  store float %load15, ptr %arg1, align 4
+  %add16 = add i32 %phi, 1
+  %icmp17 = icmp ne i32 %phi, %arg2
+  br i1 %icmp17, label %bb7, label %bb18
+
+bb18:                                             ; preds = %bb11
+  ret void
+}
+; CHECK:Invalidating scop from further optimization as max operations exceeded