From 3909dd45f0497456a974ec6583348100df9ae3ac Mon Sep 17 00:00:00 2001
From: Fabian Stuckmann <fabian.stuckmann@amd.com>
Date: Thu, 14 Nov 2024 10:50:04 +0000
Subject: [PATCH] [loop count assumptions] by default enable loop iteration
 assumption for AIE targets

---
 clang/test/CodeGen/aie/peel-itercount.c      |  4 +---
 llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/clang/test/CodeGen/aie/peel-itercount.c b/clang/test/CodeGen/aie/peel-itercount.c
index 5c02e73f5cb4..0c70350826dc 100644
--- a/clang/test/CodeGen/aie/peel-itercount.c
+++ b/clang/test/CodeGen/aie/peel-itercount.c
@@ -14,10 +14,8 @@
 // and that the itercounts have been updated appropriately
 
 // CHECK-LABEL: loop28_37
-// CHECK: entry
-// CHECK: for.body.preheader:
 // CHECK: for.body.peel.next:
-// CHECK: for.cond.cleanup: 
+// CHECK: for.cond.cleanup:
 // CHECK: for.body:
 // CHECK: !llvm.loop !6
 // CHECK: !6 = distinct !{!6, !7, !8, !9}
diff --git a/llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp b/llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp
index e35213b584ee..1c1af967d3de 100644
--- a/llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp
+++ b/llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp
@@ -43,6 +43,7 @@
 #include "llvm/Transforms/IPO/GlobalDCE.h"
 #include "llvm/Transforms/IPO/Internalize.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/LoopIterCountAssumptions.h"
 
 using namespace llvm;
 
@@ -56,6 +57,11 @@ static cl::opt<bool>
                          cl::desc("Enable tail merging for AIE."),
                          cl::init(false), cl::Hidden);
 
+static cl::opt<bool> EnableLoopIterCountToAssumptions(
+    "enable-loop-iter-count-assumptions", cl::Hidden, cl::init(true),
+    cl::desc(
+        "Enable Conversion of Loop Iteration Count Metadata to Assumptions."));
+
 // Option to run internalize pass.
 static cl::opt<bool> InternalizeSymbols(
     "aie-internalize-symbols",
@@ -179,6 +185,15 @@ void AIEBaseTargetMachine::registerPassBuilderCallbacks(
           }
         });
   }
+
+  if (EnableLoopIterCountToAssumptions) {
+    PB.registerPipelineEarlySimplificationEPCallback(
+        [](ModulePassManager &PM, OptimizationLevel) {
+          PM.addPass(createModuleToFunctionPassAdaptor(
+              createFunctionToLoopPassAdaptor(LoopIterCountAssumptions(),
+                                              /*UseMemorySSA*/ true)));
+        });
+  }
 }
 
 AIEBasePassConfig::AIEBasePassConfig(LLVMTargetMachine &TM, PassManagerBase &PM)