Convert Loop Metadata to Asserts to help Loop rotation

clang/lib/Driver/ToolChains/AIE.cpp

@@ -163,6 +163,9 @@ void AIEToolChain::addClangTargetOptions(
 
   // Extend the max limit of the search depth in BasicAA
   CC1Args.append({"-mllvm", "-basic-aa-max-lookup-search-depth=10"});
+
+  // Enable Loop Iteration Count Assumptions
+  CC1Args.append({"-mllvm", "-enable-loop-iter-count-assumptions=true"});
 }
 
 // Avoid using newer dwarf versions, as the simulator doesn't understand newer

clang/test/CodeGen/aie/peel-itercount.c

@@ -14,8 +14,9 @@
 // and that the itercounts have been updated appropriately
 
 // CHECK-LABEL: loop28_37
-// CHECK: for.body.preheader:
-// CHECK: for.body.peel.next:
+// CHECK: for.body.peel.next6:
+// CHECK: for.cond.cleanup:
+// CHECK: for.body:
 // CHECK: !llvm.loop !6
 // CHECK: !6 = distinct !{!6, !7, !8, !9}
 // CHECK: !7 = !{!"llvm.loop.peeled.count", i32 2}

llvm/include/llvm/Transforms/Utils/LoopIterCountAssumptions.h

@@ -0,0 +1,35 @@
+//===-- LoopIterCountAssumptions.h - Add loop assumptions -------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+// (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass converts Loop Iteration Count Metadata to Assumptions which can be
+// picked up by Loop Rotate to remove Loop Guards.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_UTILS_LOOPITERCOUNTASSUMPTIONS_H
+#define LLVM_TRANSFORMS_UTILS_LOOPITERCOUNTASSUMPTIONS_H
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+
+namespace llvm {
+
+class Loop;
+/// Converts Loop Iteration Count Metadata to Assumptions.
+class LoopIterCountAssumptions
+    : public PassInfoMixin<LoopIterCountAssumptions> {
+
+public:
+  PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,
+                        LoopStandardAnalysisResults &AR, LPMUpdater &U);
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_UTILS_LOOPITERCOUNTASSUMPTIONS_H

llvm/lib/Passes/PassBuilder.cpp

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2024 Advanced Micro Devices, Inc. or its
+// affiliates
+//
 //===----------------------------------------------------------------------===//
 /// \file
 ///
@@ -274,6 +277,7 @@
 #include "llvm/Transforms/Utils/InstructionNamer.h"
 #include "llvm/Transforms/Utils/LCSSA.h"
 #include "llvm/Transforms/Utils/LibCallsShrinkWrap.h"
+#include "llvm/Transforms/Utils/LoopIterCountAssumptions.h"
 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopVersioning.h"
 #include "llvm/Transforms/Utils/LowerGlobalDtors.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -132,6 +132,7 @@
 #include "llvm/Transforms/Utils/CountVisits.h"
 #include "llvm/Transforms/Utils/InjectTLIMappings.h"
 #include "llvm/Transforms/Utils/LibCallsShrinkWrap.h"
+#include "llvm/Transforms/Utils/LoopIterCountAssumptions.h"
 #include "llvm/Transforms/Utils/Mem2Reg.h"
 #include "llvm/Transforms/Utils/MoveAutoInit.h"
 #include "llvm/Transforms/Utils/NameAnonGlobals.h"
@@ -306,6 +307,11 @@ static cl::opt<bool> UseLoopVersioningLICM(
     "enable-loop-versioning-licm", cl::init(false), cl::Hidden,
     cl::desc("Enable the experimental Loop Versioning LICM pass"));
 
+static cl::opt<bool> EnableLoopIterCountToAssumptions(
+    "enable-loop-iter-count-assumptions", cl::Hidden, cl::init(false),
+    cl::desc(
+        "Enable Conversion of Loop Iteration Count Metadata to Assumptions."));
+
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
 
@@ -463,6 +469,9 @@ PassBuilder::buildO1FunctionSimplificationPipeline(OptimizationLevel Level,
   LPM1.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap,
                         /*AllowSpeculation=*/false));
 
+  if (EnableLoopIterCountToAssumptions)
+    LPM1.addPass(LoopIterCountAssumptions());
+
   LPM1.addPass(LoopRotatePass(/* Disable header duplication */ true,
                               isLTOPreLink(Phase)));
   // TODO: Investigate promotion cap for O1.
@@ -644,6 +653,9 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
   LPM1.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap,
                         /*AllowSpeculation=*/false));
 
+  if (EnableLoopIterCountToAssumptions)
+    LPM1.addPass(LoopIterCountAssumptions());
+
   // Disable header duplication in loop rotation at -Oz.
   LPM1.addPass(LoopRotatePass(EnableLoopHeaderDuplication ||
                                   Level != OptimizationLevel::Oz,

llvm/lib/Passes/PassRegistry.def

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2024 Advanced Micro Devices, Inc. or its
+// affiliates
+//
 //===----------------------------------------------------------------------===//
 //
 // This file is used as the registry of passes that are part of the core LLVM
@@ -604,6 +607,7 @@ LOOP_PASS("loop-bound-split", LoopBoundSplitPass())
 LOOP_PASS("loop-deletion", LoopDeletionPass())
 LOOP_PASS("loop-idiom", LoopIdiomRecognizePass())
 LOOP_PASS("loop-instsimplify", LoopInstSimplifyPass())
+LOOP_PASS("loop-iter-count-assumptions", LoopIterCountAssumptions())
 LOOP_PASS("loop-predication", LoopPredicationPass())
 LOOP_PASS("loop-reduce", LoopStrengthReducePass())
 LOOP_PASS("loop-simplifycfg", LoopSimplifyCFGPass())

llvm/lib/Transforms/Utils/CMakeLists.txt

@@ -1,3 +1,11 @@
+#
+# This file is licensed 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
+#
+# Modifications (c) Copyright 2024 Advanced Micro Devices, Inc. or its
+# affiliates
+#
 add_llvm_component_library(LLVMTransformUtils
   AddDiscriminators.cpp
   AMDGPUEmitPrintf.cpp
@@ -39,6 +47,7 @@ add_llvm_component_library(LLVMTransformUtils
   LibCallsShrinkWrap.cpp
   Local.cpp
   LoopConstrainer.cpp
+  LoopIterCountAssumptions.cpp
   LoopPeel.cpp
   LoopRotationUtils.cpp
   LoopSimplify.cpp