[Mlir-commits] [mlir] [MLIR][LICM] Hoist loop-invariant load operations using alias analysis. (PR #193166)
Ming Yan
llvmlistbot at llvm.org
Sun Apr 26 17:51:49 PDT 2026
https://github.com/NexMing updated https://github.com/llvm/llvm-project/pull/193166
>From 7cb6a063af556c365d62119f32e9f743e9d797e0 Mon Sep 17 00:00:00 2001
From: yanming <ming.yan at terapines.com>
Date: Tue, 21 Apr 2026 17:26:01 +0800
Subject: [PATCH] [MLIR] Hoist loop-invariant load operations using alias
analysis.
---
.../Transforms/LoopInvariantCodeMotionUtils.h | 6 +
mlir/include/mlir/Transforms/Passes.td | 5 +
.../Transforms/LoopInvariantCodeMotion.cpp | 16 ++-
.../Utils/LoopInvariantCodeMotionUtils.cpp | 73 ++++++++++++
.../loop-invariant-code-motion.mlir | 110 +++++++++++++++++-
5 files changed, 207 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index 3ceef44d799e8..6601ff591ae81 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -9,6 +9,7 @@
#ifndef MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
#define MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/SmallVector.h"
@@ -72,6 +73,11 @@ size_t moveLoopInvariantCode(
/// methods provided by the interface.
size_t moveLoopInvariantCode(LoopLikeOpInterface loopLike);
+/// Hoist loop-invariant load ops from the given loop-like op. Return the number
+/// of load ops that are hoisted.
+size_t hoistLoopInvariantLoadOps(LoopLikeOpInterface loopLike,
+ AliasAnalysis &aa);
+
/// Hoist loop-invariant tensor subsets (subset extraction and subset insertion
/// ops) from loop-like ops. Extraction ops are moved before the loop. Insertion
/// ops are moved after the loop. The loop body operates on newly added region
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 1474e580cfc03..d86c15de296ef 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -357,6 +357,11 @@ def LocationSnapshot : Pass<"snapshot-op-locations"> {
def LoopInvariantCodeMotionPass : Pass<"loop-invariant-code-motion"> {
let summary = "Hoist loop invariant instructions outside of the loop";
+ let options = [
+ Option<"useAliasAnalysis", "use-aa", "bool",
+ /*default=*/"false",
+ "Use alias analysis to hoist memory operations (experimental).">,
+ ];
}
def LoopInvariantSubsetHoistingPass : Pass<"loop-invariant-subset-hoisting"> {
diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
index 609d14d3b5638..8a35fd5ae4640 100644
--- a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
@@ -10,6 +10,7 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/IR/PatternMatch.h"
@@ -28,6 +29,8 @@ namespace {
/// Loop invariant code motion (LICM) pass.
struct LoopInvariantCodeMotion
: public impl::LoopInvariantCodeMotionPassBase<LoopInvariantCodeMotion> {
+ using impl::LoopInvariantCodeMotionPassBase<
+ LoopInvariantCodeMotion>::LoopInvariantCodeMotionPassBase;
void runOnOperation() override;
};
@@ -39,11 +42,20 @@ struct LoopInvariantSubsetHoisting
} // namespace
void LoopInvariantCodeMotion::runOnOperation() {
+ AliasAnalysis *aa = nullptr;
+ if (useAliasAnalysis) {
+ aa = &getAnalysis<AliasAnalysis>();
+ }
+
// Walk through all loops in a function in innermost-loop-first order. This
// way, we first LICM from the inner loop, and place the ops in
// the outer loop, which in turn can be further LICM'ed.
- getOperation()->walk(
- [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
+ getOperation()->walk([&](LoopLikeOpInterface loopLike) {
+ moveLoopInvariantCode(loopLike);
+ if (aa) {
+ hoistLoopInvariantLoadOps(loopLike, *aa);
+ }
+ });
}
void LoopInvariantSubsetHoisting::runOnOperation() {
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index aef81bf2e15a3..3f03131da017a 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -12,6 +12,7 @@
#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
@@ -118,6 +119,78 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
+size_t mlir::hoistLoopInvariantLoadOps(LoopLikeOpInterface loopLike,
+ AliasAnalysis &aa) {
+ // Only perform load hoisting if we have alias analysis and the loop is
+ // guaranteed to execute at least once.
+ //
+ // TODO: We can still perform load hoisting for loops with unknown trip
+ // count, but we need to be more careful about the legality of hoisting a
+ // load out of the loop.
+ std::optional<APInt> tripCount = loopLike.getStaticTripCount();
+ if (!tripCount.has_value() || *tripCount == 0)
+ return 0;
+
+ size_t numMoved = 0;
+ SmallVector<Operation *> loadOps;
+ DenseSet<MemoryEffects::EffectInstance *> writeEffects;
+
+ auto noAlias = [&](Value val1, Value val2) -> bool {
+ return aa.alias(val1, val2).isNo();
+ };
+
+ for (Region *region : loopLike.getLoopRegions())
+ for (Operation &op : region->getOps()) {
+ // Collect loop-invariant load ops. We will check that the loaded value
+ // is not written to by the loop before hoisting the load.
+ //
+ // TODO: we can also consider hoisting load ops that have multiple
+ // effects.
+ if (hasSingleEffect<MemoryEffects::Read>(&op)) {
+ loadOps.push_back(&op);
+ continue;
+ }
+
+ if (auto effects = getEffectsRecursively(&op)) {
+ // Collect all write effects in the loop. We will check that the loaded
+ // value is not written to by any of these effects before hoisting the
+ // load.
+ for (auto &effect : *effects)
+ if (isa<MemoryEffects::Write>(effect.getEffect()))
+ writeEffects.insert(&effect);
+ }
+ }
+
+ for (Operation *loadOp : loadOps) {
+ LDBG() << "Checking load op: "
+ << OpWithFlags(loadOp, OpPrintingFlags().skipRegions());
+
+ if (!canBeHoisted(loadOp, [&](Value value) {
+ return loopLike.isDefinedOutsideOfLoop(value);
+ }))
+ continue;
+
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ cast<MemoryEffectOpInterface>(loadOp).getEffects(effects);
+ assert(effects.size() == 1 &&
+ isa<MemoryEffects::Read>(effects[0].getEffect()) &&
+ "expected a single read effect");
+
+ Value loadedValue = effects[0].getValue();
+ if (llvm::all_of(writeEffects,
+ [&](MemoryEffects::EffectInstance *writeEffect) {
+ return noAlias(loadedValue, writeEffect->getValue());
+ })) {
+ LDBG() << "Hoisting loop-invariant load op: "
+ << OpWithFlags(loadOp, OpPrintingFlags().skipRegions());
+ loopLike.moveOutOfLoop(loadOp);
+ ++numMoved;
+ }
+ }
+
+ return numMoved;
+}
+
namespace {
/// Helper data structure that keeps track of equivalent/disjoint subset ops.
class MatchingSubsets {
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 31a4f64dd7de0..0351c731c84ba 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1,4 +1,5 @@
-// RUN: mlir-opt %s -split-input-file -loop-invariant-code-motion | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -loop-invariant-code-motion | FileCheck %s --check-prefixes=CHECK,NOAA
+// RUN: mlir-opt %s -split-input-file -loop-invariant-code-motion="use-aa=true" | FileCheck %s --check-prefixes=CHECK,USEAA
func.func @nested_loops_both_having_invariant_code() {
%m = memref.alloc() : memref<10xf32>
@@ -1582,3 +1583,110 @@ func.func @do_not_hoist_vector_transfer_ops_memref(
}
func.return %final : vector<4x4xf32>
}
+
+// -----
+
+// CHECK-LABEL: func @do_not_hoist_load_ops_mayalias
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: memref.load
+// CHECK: memref.load
+// CHECK: memref.load
+// CHECK: memref.store
+func.func @do_not_hoist_load_ops_mayalias(%A : memref<128x128xf32>,
+ %B : memref<128x128xf32>,
+ %C : memref<128x128xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ scf.for %i = %c0 to %c128 step %c1 {
+ scf.for %k = %c0 to %c128 step %c1 {
+ scf.for %j = %c0 to %c128 step %c1 {
+ %a_ik = memref.load %A[%i, %k] : memref<128x128xf32>
+ %b_kj = memref.load %B[%k, %j] : memref<128x128xf32>
+ %c_ij = memref.load %C[%i, %j] : memref<128x128xf32>
+ %prod = arith.mulf %a_ik, %b_kj : f32
+ %sum = arith.addf %c_ij, %prod : f32
+ memref.store %sum, %C[%i, %j] : memref<128x128xf32>
+ }
+ }
+ }
+ func.return
+}
+
+// -----
+
+// NOAA-LABEL: func @hoist_load_ops_noalias
+// NOAA: scf.for
+// NOAA: scf.for
+// NOAA: scf.for
+// NOAA: memref.load
+// NOAA: memref.load
+// NOAA: memref.load
+// NOAA: memref.store
+
+// USEAA-LABEL: func @hoist_load_ops_noalias
+// USEAA: scf.for
+// USEAA: scf.for
+// USEAA: memref.load
+// USEAA: scf.for
+// USEAA: memref.load
+// USEAA: memref.load
+// USEAA: memref.store
+func.func @hoist_load_ops_noalias(%A : memref<128x128xf32>,
+ %B : memref<128x128xf32>,
+ %C : memref<128x128xf32>) {
+ %AA, %CC = memref.distinct_objects %A, %C : memref<128x128xf32>, memref<128x128xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ scf.for %i = %c0 to %c128 step %c1 {
+ scf.for %k = %c0 to %c128 step %c1 {
+ scf.for %j = %c0 to %c128 step %c1 {
+ %a_ik = memref.load %AA[%i, %k] : memref<128x128xf32>
+ %b_kj = memref.load %B[%k, %j] : memref<128x128xf32>
+ %c_ij = memref.load %CC[%i, %j] : memref<128x128xf32>
+ %prod = arith.mulf %a_ik, %b_kj : f32
+ %sum = arith.addf %c_ij, %prod : f32
+ memref.store %sum, %CC[%i, %j] : memref<128x128xf32>
+ }
+ }
+ }
+ func.return
+}
+
+
+// -----
+
+// CHECK-LABEL: func @do_not_hoist_load_ops_unknown_tripcount
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: memref.load
+// CHECK: memref.load
+// CHECK: memref.load
+// CHECK: memref.store
+func.func @do_not_hoist_load_ops_unknown_tripcount(%A : memref<?x?xf32>,
+ %B : memref<?x?xf32>,
+ %C : memref<?x?xf32>) {
+ %AA, %CC = memref.distinct_objects %A, %C : memref<?x?xf32>, memref<?x?xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %M = memref.dim %A, %c0 : memref<?x?xf32>
+ %K = memref.dim %A, %c1 : memref<?x?xf32>
+ %N = memref.dim %B, %c1 : memref<?x?xf32>
+ scf.for %i = %c0 to %M step %c1 {
+ scf.for %k = %c0 to %K step %c1 {
+ scf.for %j = %c0 to %N step %c1 {
+ %a_ik = memref.load %AA[%i, %k] : memref<?x?xf32>
+ %b_kj = memref.load %B[%k, %j] : memref<?x?xf32>
+ %c_ij = memref.load %CC[%i, %j] : memref<?x?xf32>
+ %prod = arith.mulf %a_ik, %b_kj : f32
+ %sum = arith.addf %c_ij, %prod : f32
+ memref.store %sum, %CC[%i, %j] : memref<?x?xf32>
+ }
+ }
+ }
+ func.return
+}
More information about the Mlir-commits
mailing list