[Mlir-commits] [mlir] [mlir] Add normalize pass to MLIR (PR #186647)
lonely eagle
llvmlistbot at llvm.org
Mon Jul 6 01:44:31 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/186647
>From d6f66c9c22d70eb3722b8e520331e872234467af Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Fri, 13 Mar 2026 02:26:50 +0000
Subject: [PATCH 1/4] complete basic normalize.
---
mlir/include/mlir/Transforms/Passes.h | 1 +
mlir/include/mlir/Transforms/Passes.td | 8 ++
mlir/lib/Transforms/CMakeLists.txt | 1 +
mlir/lib/Transforms/Normalize.cpp | 112 +++++++++++++++++++++++++
mlir/test/Transforms/normalize.mlir | 84 +++++++++++++++++++
5 files changed, 206 insertions(+)
create mode 100644 mlir/lib/Transforms/Normalize.cpp
create mode 100644 mlir/test/Transforms/normalize.mlir
diff --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index 313755fde3479..be56a548f63b1 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -42,6 +42,7 @@ class GreedyRewriteConfig;
#define GEN_PASS_DECL_LOOPINVARIANTSUBSETHOISTINGPASS
#define GEN_PASS_DECL_INLINERPASS
#define GEN_PASS_DECL_MEM2REG
+#define GEN_PASS_DECL_NORMALIZEPASS
#define GEN_PASS_DECL_PRINTIRPASS
#define GEN_PASS_DECL_PRINTOPSTATSPASS
#define GEN_PASS_DECL_REMOVEDEADVALUESPASS
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 74ac370ea950b..10193bbcef2bd 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -656,4 +656,12 @@ def BubbleDownMemorySpaceCasts :
}];
}
+def NormalizePass : InterfacePass<"normalize", "FunctionOpInterface"> {
+ let summary = "Transforms IR into a normal form that's easier to diff.";
+ let description = [{
+ This pass attempts to relocate the defining ops of operands for any
+ side-effecting or terminator operation to their nearest dominating positions.
+ }];
+}
+
#endif // MLIR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index 66b39f53c91df..41347d23b0f49 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -12,6 +12,7 @@ add_mlir_library(MLIRTransforms
LocationSnapshot.cpp
LoopInvariantCodeMotion.cpp
Mem2Reg.cpp
+ Normalize.cpp
OpStats.cpp
PrintIR.cpp
RemoveDeadValues.cpp
diff --git a/mlir/lib/Transforms/Normalize.cpp b/mlir/lib/Transforms/Normalize.cpp
new file mode 100644
index 0000000000000..1904d25d36674
--- /dev/null
+++ b/mlir/lib/Transforms/Normalize.cpp
@@ -0,0 +1,112 @@
+//===- Normalize.cpp - Transforms IR into a normal form ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/DebugLog.h"
+
+using namespace mlir;
+
+namespace mlir {
+#define GEN_PASS_DEF_NORMALIZEPASS
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "normalize"
+
+namespace {
+
+bool isOutput(Operation *op) {
+ if (!op)
+ return false;
+ return !isMemoryEffectFree(op) || op->hasTrait<OpTrait::IsTerminator>();
+}
+
+/// Returns a vector of output ops. An output is a op which
+/// has side-effects or is terminator.
+SmallVector<Operation *> collectOutputs(Operation *root) {
+ SmallVector<Operation *> outputs;
+ root->walk([&](Operation *op) {
+ if (isOutput(op))
+ outputs.push_back(op);
+ });
+ return outputs;
+}
+
+/// The function returns the operation that dominates all other operations in
+/// the given list.
+Operation *getDominateOp(SmallVectorImpl<Operation *> &ops) {
+ if (ops.empty())
+ return {};
+ Operation *curDomOp = ops.front();
+ DominanceInfo domInfo(curDomOp);
+ for (size_t i = 1, e = ops.size(); i < e; ++i) {
+ bool dominateA = domInfo.dominates(ops[i], curDomOp);
+ bool dominateB = domInfo.dominates(curDomOp, ops[i]);
+ if (dominateA) {
+ LDBG() << OpWithFlags(ops[i], OpPrintingFlags().skipRegions())
+ << "\ndominate\n"
+ << OpWithFlags(curDomOp, OpPrintingFlags().skipRegions());
+ curDomOp = ops[i];
+ }
+ if (!dominateA && !dominateB) {
+ LDBG() << OpWithFlags(ops[i], OpPrintingFlags().skipRegions())
+ << "\nand\n"
+ << OpWithFlags(curDomOp, OpPrintingFlags().skipRegions())
+ << "\ndo not dominate each other";
+ return {};
+ }
+ }
+ return curDomOp;
+}
+
+/// Move used to its nearest user and recursively perform the same process on
+/// the defining operations of its operands.
+void reorderOutput(IRRewriter &rewriter, Operation *used) {
+ if (!isPure(used))
+ return;
+ SmallVector<Operation *> users(used->getUsers());
+ if (Operation *domOp = getDominateOp(users)) {
+ rewriter.moveOpBefore(used, domOp);
+ for (Value operand : used->getOperands())
+ if (Operation *defineOp = operand.getDefiningOp())
+ reorderOutput(rewriter, defineOp);
+ }
+}
+
+/// Reorders ops by walking up the tree from each operand of an output op and
+/// reducing the def-use distance. This method assumes that output ops were
+/// collected top-down, otherwise the def-use chain may be broken. This method
+/// is a wrapper for recursive reorderOutput().
+void reorderOutputs(IRRewriter &rewriter,
+ SmallVectorImpl<Operation *> &outputs) {
+ SmallPtrSet<Operation *, 16> visited;
+ for (Operation *output : outputs) {
+ for (Value operand : output->getOperands()) {
+ if (Operation *defineOp = operand.getDefiningOp();
+ defineOp && !visited.contains(defineOp)) {
+ reorderOutput(rewriter, defineOp);
+ }
+ }
+ }
+}
+
+struct NormalizePass : public impl::NormalizePassBase<NormalizePass> {
+ using impl::NormalizePassBase<NormalizePass>::NormalizePassBase;
+ void runOnOperation() override;
+};
+} // namespace
+
+void NormalizePass::runOnOperation() {
+ IRRewriter rewriter(&getContext());
+ SmallVector<Operation *> outputs = collectOutputs(getOperation());
+ reorderOutputs(rewriter, outputs);
+}
diff --git a/mlir/test/Transforms/normalize.mlir b/mlir/test/Transforms/normalize.mlir
new file mode 100644
index 0000000000000..241e28c3c385d
--- /dev/null
+++ b/mlir/test/Transforms/normalize.mlir
@@ -0,0 +1,84 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(normalize))" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @multiple_memref_store
+// CHECK-SAME: %[[ARG0:.*]]: index,
+// CHECK-SAME: %[[ARG1:.*]]: memref<?xf32>
+func.func @multiple_memref_store(%arg0: index, %arg1 : memref<?xf32>) {
+ %f0 = arith.constant 0.0 : f32
+ %f1 = arith.constant 1.0 : f32
+ %add = arith.addi %arg0, %arg0 : index
+ %sub = arith.subi %arg0, %arg0 : index
+ memref.store %f0, %arg1[%add] : memref<?xf32>
+ memref.store %f1, %arg1[%sub] : memref<?xf32>
+ return
+}
+
+// CHECK-NEXT: %[[C0:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-NEXT: %[[ADD:.*]] = arith.addi %[[ARG0]], %[[ARG0]] : index
+// CHECK-NEXT: memref.store %[[C0]], %[[ARG1]]{{\[}}%[[ADD]]] : memref<?xf32>
+// CHECK-NEXT: %[[C1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-NEXT: %[[SUB:.*]] = arith.subi %[[ARG0]], %[[ARG0]] : index
+// CHECK-NEXT: memref.store %[[C1]], %[[ARG1]]{{\[}}%[[SUB]]] : memref<?xf32>
+
+// -----
+
+// CHECK-LABEL: func @normalize_return
+// CHECK-SAME: %[[ARG0:.*]]: index,
+// CHECK-SAME: %[[ARG1:.*]]: memref<?xf32>
+func.func @normalize_return(%arg0: index, %arg1 : memref<?xf32>) -> index {
+ %f0 = arith.constant 0.0 : f32
+ %add = arith.addi %arg0, %arg0 : index
+ %sub = arith.subi %arg0, %arg0 : index
+ memref.store %f0, %arg1[%add] : memref<?xf32>
+ return %sub : index
+}
+
+// CHECK: memref.store
+// CHECK-NEXT: %[[SUB:.*]] = arith.subi %[[ARG0]], %[[ARG0]] : index
+// CHECK-NEXT: return %[[SUB]] : index
+
+// -----
+
+// CHECK-LABEL: func @cross_region
+// CHECK-SAME: %[[ARG0:.*]]: f32,
+// CHECK-SAME: %[[ARG1:.*]]: memref<10xf32>
+func.func @cross_region(%arg0: f32, %arg1 : memref<10xf32>) {
+ %add = arith.addf %arg0, %arg0 : f32
+ affine.for %i = 0 to 5 {
+ memref.store %add, %arg1[%i] : memref<10xf32>
+ }
+ %exp = math.log2 %add : f32
+ affine.for %i = 6 to 10 {
+ memref.store %exp, %arg1[%i] : memref<10xf32>
+ }
+ return
+}
+
+// CHECK: affine.for %[[IV:.*]] = 6 to 10 {
+// CHECK-NEXT: %[[LOG:.*]] = math.log2
+// CHECK-NEXT: memref.store %[[LOG]], %[[ARG1]]{{\[}}%[[IV]]] : memref<10xf32>
+// CHECK-NEXT: }
+
+// -----
+
+
+// CHECK-LABEL: func @side_effect_for_op
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
+func.func @side_effect_for_op(%arg1 : memref<?xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %upper = memref.dim %arg1, %c0 : memref<?xf32>
+ %f1 = arith.constant 1.0 : f32
+ scf.for %i = %c0 to %upper step %c1 {
+ memref.store %f1, %arg1[%i] : memref<?xf32>
+ }
+ return
+}
+
+// CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG0]], %[[C0]] : memref<?xf32>
+// CHECK-NEXT: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT: scf.for %[[IV:.*]] = %[[C0]] to %[[DIM]] step %[[C1]] {
+// CHECK-NEXT: %[[F1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-NEXT: memref.store %[[F1]], %[[ARG0]]{{\[}}%[[IV]]] : memref<?xf32>
+// CHECK-NEXT: }
>From b025af24e07486c0142bec3f76c0e9cafb559d3f Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 16 Mar 2026 06:58:47 +0000
Subject: [PATCH 2/4] fix nit and add more comment.
---
mlir/include/mlir/Transforms/Passes.td | 10 ++++++++--
mlir/lib/Transforms/Normalize.cpp | 20 ++++++++++---------
mlir/test/Transforms/normalize.mlir | 27 +++++++++++++-------------
3 files changed, 33 insertions(+), 24 deletions(-)
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 10193bbcef2bd..a03c69c12461e 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -659,8 +659,14 @@ def BubbleDownMemorySpaceCasts :
def NormalizePass : InterfacePass<"normalize", "FunctionOpInterface"> {
let summary = "Transforms IR into a normal form that's easier to diff.";
let description = [{
- This pass attempts to relocate the defining ops of operands for any
- side-effecting or terminator operation to their nearest dominating positions.
+ This pass aims to transform MLIR Modules into a normal form by reordering
+ operations while preserving the same semantics. It attempts to relocate
+ the defining ops of operands for any side-effecting or terminator operation
+ to their nearest dominating positions.
+
+ Note: The pass trying to increase syntactic equivalence of code to reduce
+ diff size while retaining semantic equivalence. It cannot replace the
+ canonicalization pass (a pass that aims for increasing semantic equivalence instead).
}];
}
diff --git a/mlir/lib/Transforms/Normalize.cpp b/mlir/lib/Transforms/Normalize.cpp
index 1904d25d36674..7e63c42b6d163 100644
--- a/mlir/lib/Transforms/Normalize.cpp
+++ b/mlir/lib/Transforms/Normalize.cpp
@@ -43,21 +43,22 @@ SmallVector<Operation *> collectOutputs(Operation *root) {
/// The function returns the operation that dominates all other operations in
/// the given list.
-Operation *getDominateOp(SmallVectorImpl<Operation *> &ops) {
+Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops) {
if (ops.empty())
return {};
Operation *curDomOp = ops.front();
DominanceInfo domInfo(curDomOp);
for (size_t i = 1, e = ops.size(); i < e; ++i) {
bool dominateA = domInfo.dominates(ops[i], curDomOp);
- bool dominateB = domInfo.dominates(curDomOp, ops[i]);
if (dominateA) {
LDBG() << OpWithFlags(ops[i], OpPrintingFlags().skipRegions())
<< "\ndominate\n"
<< OpWithFlags(curDomOp, OpPrintingFlags().skipRegions());
curDomOp = ops[i];
+ continue;
}
- if (!dominateA && !dominateB) {
+ bool dominateB = domInfo.dominates(curDomOp, ops[i]);
+ if (!dominateB) {
LDBG() << OpWithFlags(ops[i], OpPrintingFlags().skipRegions())
<< "\nand\n"
<< OpWithFlags(curDomOp, OpPrintingFlags().skipRegions())
@@ -70,13 +71,13 @@ Operation *getDominateOp(SmallVectorImpl<Operation *> &ops) {
/// Move used to its nearest user and recursively perform the same process on
/// the defining operations of its operands.
-void reorderOutput(IRRewriter &rewriter, Operation *used) {
- if (!isPure(used))
+void reorderOutput(IRRewriter &rewriter, Operation *producer) {
+ if (!isPure(producer))
return;
- SmallVector<Operation *> users(used->getUsers());
+ SmallVector<Operation *> users(producer->getUsers());
if (Operation *domOp = getDominateOp(users)) {
- rewriter.moveOpBefore(used, domOp);
- for (Value operand : used->getOperands())
+ rewriter.moveOpBefore(producer, domOp);
+ for (Value operand : producer->getOperands())
if (Operation *defineOp = operand.getDefiningOp())
reorderOutput(rewriter, defineOp);
}
@@ -87,12 +88,13 @@ void reorderOutput(IRRewriter &rewriter, Operation *used) {
/// collected top-down, otherwise the def-use chain may be broken. This method
/// is a wrapper for recursive reorderOutput().
void reorderOutputs(IRRewriter &rewriter,
- SmallVectorImpl<Operation *> &outputs) {
+ const SmallVectorImpl<Operation *> &outputs) {
SmallPtrSet<Operation *, 16> visited;
for (Operation *output : outputs) {
for (Value operand : output->getOperands()) {
if (Operation *defineOp = operand.getDefiningOp();
defineOp && !visited.contains(defineOp)) {
+ visited.insert(defineOp);
reorderOutput(rewriter, defineOp);
}
}
diff --git a/mlir/test/Transforms/normalize.mlir b/mlir/test/Transforms/normalize.mlir
index 241e28c3c385d..28f0c05263278 100644
--- a/mlir/test/Transforms/normalize.mlir
+++ b/mlir/test/Transforms/normalize.mlir
@@ -22,20 +22,17 @@ func.func @multiple_memref_store(%arg0: index, %arg1 : memref<?xf32>) {
// -----
-// CHECK-LABEL: func @normalize_return
-// CHECK-SAME: %[[ARG0:.*]]: index,
-// CHECK-SAME: %[[ARG1:.*]]: memref<?xf32>
-func.func @normalize_return(%arg0: index, %arg1 : memref<?xf32>) -> index {
- %f0 = arith.constant 0.0 : f32
- %add = arith.addi %arg0, %arg0 : index
- %sub = arith.subi %arg0, %arg0 : index
- memref.store %f0, %arg1[%add] : memref<?xf32>
- return %sub : index
+// CHECK-LABEL: func @return_multiple_operands
+// CHECK-SAME: %[[ARG0:.*]]: index
+func.func @return_multiple_operands (%arg0: index) -> (index, index) {
+ %0 = arith.addi %arg0, %arg0 : index
+ %1 = arith.subi %arg0, %arg0 : index
+ return %1, %0 : index, index
}
-// CHECK: memref.store
// CHECK-NEXT: %[[SUB:.*]] = arith.subi %[[ARG0]], %[[ARG0]] : index
-// CHECK-NEXT: return %[[SUB]] : index
+// CHECK-NEXT: %[[ADD:.*]] = arith.addi %[[ARG0]], %[[ARG0]] : index
+// CHECK-NEXT: return %[[SUB]], %[[ADD]] : index, index
// -----
@@ -61,10 +58,14 @@ func.func @cross_region(%arg0: f32, %arg1 : memref<10xf32>) {
// -----
+// This test verifies the reordering of scf.for ops.
+// The memref.store within the scf.for causes the loop to have side effects.
+// The lower bound of the scf.for remains in its original position
+// because the upper bound depends on it, but the step has been reordered.
-// CHECK-LABEL: func @side_effect_for_op
+// CHECK-LABEL: func @side_effect_loop_op
// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
-func.func @side_effect_for_op(%arg1 : memref<?xf32>) {
+func.func @side_effect_loop_op(%arg1 : memref<?xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%upper = memref.dim %arg1, %c0 : memref<?xf32>
>From c91b6a5f4b2bada8263b6063854c8cb39fc408c7 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 19 Apr 2026 05:00:23 +0000
Subject: [PATCH 3/4] use getAnalysis to get DominanceInfo and use
markAnalysesPreserved.
---
mlir/lib/Transforms/Normalize.cpp | 22 ++++++++++++++--------
1 file changed, 14 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Transforms/Normalize.cpp b/mlir/lib/Transforms/Normalize.cpp
index 7e63c42b6d163..6f42c61ee7d18 100644
--- a/mlir/lib/Transforms/Normalize.cpp
+++ b/mlir/lib/Transforms/Normalize.cpp
@@ -43,11 +43,11 @@ SmallVector<Operation *> collectOutputs(Operation *root) {
/// The function returns the operation that dominates all other operations in
/// the given list.
-Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops) {
+Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops,
+ const DominanceInfo &domInfo) {
if (ops.empty())
return {};
Operation *curDomOp = ops.front();
- DominanceInfo domInfo(curDomOp);
for (size_t i = 1, e = ops.size(); i < e; ++i) {
bool dominateA = domInfo.dominates(ops[i], curDomOp);
if (dominateA) {
@@ -71,15 +71,16 @@ Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops) {
/// Move used to its nearest user and recursively perform the same process on
/// the defining operations of its operands.
-void reorderOutput(IRRewriter &rewriter, Operation *producer) {
+void reorderOutput(IRRewriter &rewriter, Operation *producer,
+ const DominanceInfo &domInfo) {
if (!isPure(producer))
return;
SmallVector<Operation *> users(producer->getUsers());
- if (Operation *domOp = getDominateOp(users)) {
+ if (Operation *domOp = getDominateOp(users, domInfo)) {
rewriter.moveOpBefore(producer, domOp);
for (Value operand : producer->getOperands())
if (Operation *defineOp = operand.getDefiningOp())
- reorderOutput(rewriter, defineOp);
+ reorderOutput(rewriter, defineOp, domInfo);
}
}
@@ -88,14 +89,15 @@ void reorderOutput(IRRewriter &rewriter, Operation *producer) {
/// collected top-down, otherwise the def-use chain may be broken. This method
/// is a wrapper for recursive reorderOutput().
void reorderOutputs(IRRewriter &rewriter,
- const SmallVectorImpl<Operation *> &outputs) {
+ const SmallVectorImpl<Operation *> &outputs,
+ const DominanceInfo &domInfo) {
SmallPtrSet<Operation *, 16> visited;
for (Operation *output : outputs) {
for (Value operand : output->getOperands()) {
if (Operation *defineOp = operand.getDefiningOp();
defineOp && !visited.contains(defineOp)) {
visited.insert(defineOp);
- reorderOutput(rewriter, defineOp);
+ reorderOutput(rewriter, defineOp, domInfo);
}
}
}
@@ -108,7 +110,11 @@ struct NormalizePass : public impl::NormalizePassBase<NormalizePass> {
} // namespace
void NormalizePass::runOnOperation() {
+ DominanceInfo &domInfo = getAnalysis<DominanceInfo>();
IRRewriter rewriter(&getContext());
SmallVector<Operation *> outputs = collectOutputs(getOperation());
- reorderOutputs(rewriter, outputs);
+ reorderOutputs(rewriter, outputs, domInfo);
+ // Since we only changed the positions of the operations, DominanceInfo and
+ // PostDominanceInfo are marked as preserved.
+ markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
}
>From 6ec1af226d1cebf159ccfd1c17406078f5b772fa Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 6 Jul 2026 08:44:03 +0000
Subject: [PATCH 4/4] update code.
---
mlir/lib/Transforms/Normalize.cpp | 365 +++++++++++++++++++++++++++---
1 file changed, 338 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Transforms/Normalize.cpp b/mlir/lib/Transforms/Normalize.cpp
index 6f42c61ee7d18..5de76bd681dd9 100644
--- a/mlir/lib/Transforms/Normalize.cpp
+++ b/mlir/lib/Transforms/Normalize.cpp
@@ -6,12 +6,29 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Block.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Region.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/Passes.h"
+#include "mlir/Support/WalkResult.h"
+#include "mlir/Transforms/CommutativityUtils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
#include "llvm/Support/DebugLog.h"
+#include <cstdint>
using namespace mlir;
@@ -24,25 +41,104 @@ namespace mlir {
namespace {
-bool isOutput(Operation *op) {
+class Normalize {
+public:
+ Normalize(IRRewriter &rewriter, DominanceInfo &domInfo)
+ : rewriter(rewriter), domInfo(domInfo) {}
+ /// Collect a vector of output ops within in \p root.
+ void collectOutputs(Operation *root);
+
+ /// Reorders ops by walking up the tree from each operand of an output op and
+ /// reducing the def-use distance. This method assumes that output ops were
+ /// collected top-down, otherwise the def-use chain may be broken. This method
+ /// is a wrapper for recursive reorderOutput().
+ void reorderOutputs();
+
+ /// Assigns unique, sequential names (e.g., "a0", "a1") to all block arguments
+ /// within \p root.
+ void nameBlockArguments(Operation *root);
+
+ /// Assigns unique, sequential names to all collected output operations.
+ void nameOperations();
+
+ /// Fold the operation name within \p root.
+ void foldOperationsName(Operation *root);
+
+ /// Greedily applies commutativity patterns using \p root to define the
+ /// transformation scope.
+ LogicalResult sortCommutativeOperands(Operation *root);
+
+private:
+ /// Reorders operations along the def-use chain from left to right, bottom to
+ /// top, starting from \p producer.
+ void reorderOutput(Operation *producer);
+
+ /// Returns true if the \p is a terminator or contains memory/side effects.
+ bool isOutput(Operation *op);
+
+ /// Returns true if the \p op is an initial operation (has no operands or only
+ /// constant-like operands).
+ bool isInitialOpeartion(Operation *op);
+ void nameOpeartion(Operation *op, SmallPtrSet<Operation *, 32> &visited);
+ void nameAsInitialOpeartion(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited);
+
+ void nameAsRegularOpeartion(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited);
+ SetVector<int> getOutputFootprint(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited);
+ void foldOperationName(Operation *op);
+ IRRewriter &rewriter;
+ DominanceInfo &domInfo;
+ SmallVector<Operation *> outputs;
+ // Random constant for hashing, so the state isn't zero.
+ const uint64_t magicHashConstant = 0x6acaa36bef8325c5ULL;
+};
+
+// Frozen mixer; basic-block names derived from these hashes appear in
+// the normalized IR text and must be deterministic across processes
+// for the normalizer's "compare normalized IR" workflow to work.
+static constexpr uint64_t hash_16_bytes(uint64_t Low, uint64_t High) {
+ const uint64_t kMul = 0x9ddfea08eb382d69ULL;
+ uint64_t A = (Low ^ High) * kMul;
+ A ^= (A >> 47);
+ uint64_t B = (High ^ A) * kMul;
+ B ^= (B >> 47);
+ B *= kMul;
+ return B;
+}
+
+/// Computes the 64-bit FNV-1a hash value of the given string \p data.
+static constexpr uint64_t strHash(std::string_view data) noexcept {
+ const uint64_t fnvOffset = 0xcbf29ce484222325ULL;
+ const uint64_t fnvPrime = 0x100000001b3ULL;
+ uint64_t hash = fnvOffset;
+ for (const auto &c : data) {
+ hash ^= static_cast<uint64_t>(c);
+ hash *= fnvPrime;
+ }
+ return hash;
+}
+
+bool Normalize::isOutput(Operation *op) {
if (!op)
return false;
return !isMemoryEffectFree(op) || op->hasTrait<OpTrait::IsTerminator>();
}
-/// Returns a vector of output ops. An output is a op which
-/// has side-effects or is terminator.
-SmallVector<Operation *> collectOutputs(Operation *root) {
- SmallVector<Operation *> outputs;
+void Normalize::collectOutputs(Operation *root) {
root->walk([&](Operation *op) {
- if (isOutput(op))
+ if (op == root)
+ return WalkResult::advance();
+ if (isOutput(op)) {
+ LDBG() << "insert " << OpWithFlags(op, OpPrintingFlags().skipRegions())
+ << " to outputs";
outputs.push_back(op);
+ }
+ return WalkResult::advance();
});
- return outputs;
}
-/// The function returns the operation that dominates all other operations in
-/// the given list.
Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops,
const DominanceInfo &domInfo) {
if (ops.empty())
@@ -69,10 +165,7 @@ Operation *getDominateOp(const SmallVectorImpl<Operation *> &ops,
return curDomOp;
}
-/// Move used to its nearest user and recursively perform the same process on
-/// the defining operations of its operands.
-void reorderOutput(IRRewriter &rewriter, Operation *producer,
- const DominanceInfo &domInfo) {
+void Normalize::reorderOutput(Operation *producer) {
if (!isPure(producer))
return;
SmallVector<Operation *> users(producer->getUsers());
@@ -80,27 +173,238 @@ void reorderOutput(IRRewriter &rewriter, Operation *producer,
rewriter.moveOpBefore(producer, domOp);
for (Value operand : producer->getOperands())
if (Operation *defineOp = operand.getDefiningOp())
- reorderOutput(rewriter, defineOp, domInfo);
+ reorderOutput(defineOp);
}
}
-/// Reorders ops by walking up the tree from each operand of an output op and
-/// reducing the def-use distance. This method assumes that output ops were
-/// collected top-down, otherwise the def-use chain may be broken. This method
-/// is a wrapper for recursive reorderOutput().
-void reorderOutputs(IRRewriter &rewriter,
- const SmallVectorImpl<Operation *> &outputs,
- const DominanceInfo &domInfo) {
+void Normalize::reorderOutputs() {
SmallPtrSet<Operation *, 16> visited;
for (Operation *output : outputs) {
for (Value operand : output->getOperands()) {
if (Operation *defineOp = operand.getDefiningOp();
defineOp && !visited.contains(defineOp)) {
visited.insert(defineOp);
- reorderOutput(rewriter, defineOp, domInfo);
+ reorderOutput(defineOp);
+ }
+ }
+ }
+}
+
+bool Normalize::isInitialOpeartion(Operation *op) {
+ for (Value operand : op->getOperands()) {
+ if (Operation *define = operand.getDefiningOp();
+ !define || !define->hasTrait<OpTrait::ConstantLike>())
+ return false;
+ }
+ return true;
+}
+
+SetVector<int>
+Normalize::getOutputFootprint(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited) {
+ SetVector<int> outputs;
+ if (visited.count(op))
+ return outputs;
+ visited.insert(op);
+
+ if (isOutput(op)) {
+ int count = 0;
+ do {
+ Region *parentRegion = op->getParentRegion();
+ int distance = 0;
+ for (Operation &curOp : parentRegion->getOps()) {
+ if (&curOp == op)
+ break;
+ ++distance;
}
+ count += distance;
+ op = parentRegion->getParentOp();
+ if (isa<FunctionOpInterface>(op))
+ break;
+ } while (!isa<FunctionOpInterface>(op));
+ outputs.insert(count);
+ return outputs;
+ }
+
+ for (Operation *user : op->getUsers()) {
+ SetVector<int> outputsUser = getOutputFootprint(user, visited);
+ outputs.insert(outputsUser.begin(), outputsUser.end());
+ }
+ return outputs;
+}
+
+void appendCallAndOperandNames(Operation *op, SmallString<512> &name,
+ SmallVectorImpl<StringRef> &operandNames) {
+ // In case of CallInst, consider callee in the instruction name.
+ if (auto callOp = dyn_cast<CallOpInterface>(op))
+ if (auto funcOp = dyn_cast<FunctionOpInterface>(callOp.resolveCallable()))
+ name.append(funcOp.getNameAttr());
+
+ if (operandNames.size() > 0) {
+ name.append("$-");
+ for (size_t i = 0, e = operandNames.size(); i < e; ++i) {
+ name.append(operandNames[i]);
+ if (i < e - 1)
+ name.append(".");
+ }
+ name.append("-$");
+ }
+ NameLoc loc = NameLoc::get(StringAttr::get(op->getContext(), name));
+ LDBG() << "set NameLoc: " << loc
+ << "\nfor: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ op->setLoc(loc);
+}
+
+void Normalize::nameAsInitialOpeartion(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited) {
+ // Operation operands for further sorting.
+ SmallVector<StringRef, 4> operandNames;
+ for (Value operand : op->getOperands()) {
+ if (Operation *define = operand.getDefiningOp())
+ nameOpeartion(define, visited);
+ if (NameLoc loc = dyn_cast<NameLoc>(operand.getLoc()))
+ operandNames.push_back(loc.getName());
+ }
+
+ if (op->use_empty())
+ return;
+
+ // Initialize to a magic constant, so the state isn't zero.
+ uint64_t hash = magicHashConstant;
+
+ // Consider instruction's opcode in the hash.
+ hash = hash_16_bytes(hash, strHash(op->getName().getStringRef().str()));
+
+ // Get output footprint for \p op.
+ SmallPtrSet<Operation *, 32> visitedOutputFoot;
+ SetVector<int> outputFootprint = getOutputFootprint(op, visitedOutputFoot);
+ LDBG() << "getOutputFootprint success";
+
+ // Consider output footprint in the hash.
+ for (const int &output : outputFootprint)
+ hash = hash_16_bytes(hash, output);
+
+ hash = hash_16_bytes(
+ hash, std::distance(op->getBlock()->begin(), op->getIterator()));
+
+ // Base operation name.
+ SmallString<512> name;
+ name.append("vl" + std::to_string(hash).substr(0, 5));
+
+ // Append call and operand name.
+ appendCallAndOperandNames(op, name, operandNames);
+}
+
+void Normalize::nameAsRegularOpeartion(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited) {
+ SmallVector<StringRef, 2> operandNames;
+ for (Value operand : op->getOperands()) {
+ if (Operation *define = operand.getDefiningOp())
+ nameOpeartion(define, visited);
+ if (NameLoc loc = dyn_cast<NameLoc>(operand.getLoc())) {
+ operandNames.push_back(loc.getName());
+ }
+ }
+
+ if (op->use_empty())
+ return;
+
+ // Initialize to a magic constant, so the state isn't zero.
+ uint64_t hash = magicHashConstant;
+
+ // Consider operation opcode in the hash.
+ uint64_t ophash = strHash(op->getName().getStringRef().str());
+ hash = hash_16_bytes(hash, ophash);
+
+ // Operand opcodes for further sorting (commutative).
+ SmallVector<int, 4> operandsOpcodes;
+
+ for (Value operand : op->getOperands())
+ if (Operation *define = operand.getDefiningOp())
+ hash =
+ hash_16_bytes(hash, strHash(define->getName().getStringRef().str()));
+
+ hash = hash_16_bytes(
+ hash, std::distance(op->getBlock()->begin(), op->getIterator()));
+
+ // Base operation name.
+ SmallString<512> name;
+ name.append("op" + std::to_string(hash).substr(0, 5));
+
+ // Append call and operand name.
+ appendCallAndOperandNames(op, name, operandNames);
+}
+
+void Normalize::nameOpeartion(Operation *op,
+ SmallPtrSet<Operation *, 32> &visited) {
+ if (visited.count(op))
+ return;
+ visited.insert(op);
+ LDBG() << "rename : " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ // Determine the type of instruction to name.
+ if (isInitialOpeartion(op)) {
+ // This is an initial instruction.
+ nameAsInitialOpeartion(op, visited);
+ } else {
+ // This must be a regular instruction.
+ nameAsRegularOpeartion(op, visited);
+ }
+}
+
+void Normalize::nameOperations() {
+ SmallPtrSet<Operation *, 32> visited;
+ for (Operation *op : outputs)
+ nameOpeartion(op, visited);
+}
+
+void Normalize::nameBlockArguments(Operation *root) {
+ size_t argumentCount = 0;
+ MLIRContext *context = root->getContext();
+ root->walk<WalkOrder::PreOrder>([&](Block *b) {
+ for (auto argument : b->getArguments()) {
+ NameLoc loc =
+ NameLoc::get(StringAttr::get(context, "a" + Twine(argumentCount++)));
+ argument.setLoc(loc);
}
+ return;
+ });
+}
+
+void Normalize::foldOperationName(Operation *op) {
+ NameLoc loc = dyn_cast<NameLoc>(op->getLoc());
+ if (!loc || loc.getName().empty() || loc.getName().str().substr(0, 2) != "op")
+ return;
+
+ SmallString<512> name;
+ name.append(loc.getName().str().substr(0, 7));
+ name.append("$-");
+ for (size_t i = 0, e = op->getNumOperands(); i < e; ++i) {
+ if (NameLoc loc = dyn_cast<NameLoc>(op->getOperand(i).getLoc()))
+ name.append(loc.getName().str().substr(0, 7));
+ if (i < e - 1)
+ name.append(".");
}
+ name.append("-$");
+
+ loc = NameLoc::get(StringAttr::get(op->getContext(), name));
+ LDBG() << "set NameLoc: " << loc
+ << "\nfor: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ op->setLoc(loc);
+}
+
+void Normalize::foldOperationsName(Operation *root) {
+ root->walk<WalkOrder::PreOrder>(
+ [&](Operation *op) { foldOperationName(op); });
+}
+
+LogicalResult Normalize::sortCommutativeOperands(Operation *root) {
+ MLIRContext *context = root->getContext();
+ RewritePatternSet patterns(context);
+ populateCommutativityUtilsPatterns(patterns);
+ if (failed(applyPatternsGreedily(root, std::move(patterns))))
+ return failure();
+ return success();
}
struct NormalizePass : public impl::NormalizePassBase<NormalizePass> {
@@ -112,9 +416,16 @@ struct NormalizePass : public impl::NormalizePassBase<NormalizePass> {
void NormalizePass::runOnOperation() {
DominanceInfo &domInfo = getAnalysis<DominanceInfo>();
IRRewriter rewriter(&getContext());
- SmallVector<Operation *> outputs = collectOutputs(getOperation());
- reorderOutputs(rewriter, outputs, domInfo);
- // Since we only changed the positions of the operations, DominanceInfo and
- // PostDominanceInfo are marked as preserved.
+ Normalize normalize(rewriter, domInfo);
+ if (failed(normalize.sortCommutativeOperands(getOperation())))
+ signalPassFailure();
+ normalize.collectOutputs(getOperation());
+ normalize.reorderOutputs();
+ normalize.nameBlockArguments(getOperation());
+ normalize.nameOperations();
+ normalize.foldOperationsName(getOperation());
+
+ // Since we only changed the positions of the operations, `DominanceInfo` and
+ // `PostDominanceInfo` are marked as preserved.
markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
}
More information about the Mlir-commits
mailing list