[Mlir-commits] [mlir] [mlir][memref] Add a folder for chained AssumeAlignmentOp ops. (PR #142425)
Han-Chung Wang
llvmlistbot at llvm.org
Mon Jun 2 09:28:16 PDT 2025
https://github.com/hanhanW created https://github.com/llvm/llvm-project/pull/142425
The chained ops can be folded away when they have the same alignment.
>From 58b37d76e8c0323d1f1c40e69806c259d244290b Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Mon, 2 Jun 2025 09:26:31 -0700
Subject: [PATCH] [mlir][memref] Add a folder for chained AssumeAlignmentOp
ops.
The chained ops can be folded away when they have the same alignment.
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
---
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td | 1 +
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 9 +++++++++
mlir/test/Dialect/MemRef/canonicalize.mlir | 13 +++++++++++++
3 files changed, 23 insertions(+)
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index f33ecb28d27cd..77e3074661abf 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -174,6 +174,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [
}];
let hasVerifier = 1;
+ let hasFolder = 1;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index aa9587510670c..d56b32193765e 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -533,6 +533,15 @@ void AssumeAlignmentOp::getAsmResultNames(
setNameFn(getResult(), "assume_align");
}
+OpFoldResult AssumeAlignmentOp::fold(FoldAdaptor adaptor) {
+ auto source = getMemref().getDefiningOp<AssumeAlignmentOp>();
+ if (!source)
+ return {};
+ if (source.getAlignment() != getAlignment())
+ return {};
+ return getMemref();
+}
+
//===----------------------------------------------------------------------===//
// CastOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index da6faf490653d..7a267ae8a2c95 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1177,3 +1177,16 @@ func.func @cannot_fold_transpose_cast(%arg0: memref<?x4xf32>) -> memref<?x?xf32,
// CHECK: return %[[TRANSPOSE]]
return %transpose : memref<?x?xf32, #transpose_map>
}
+
+// -----
+
+// CHECK-LABEL: func @fold_assume_alignment_chain
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+func.func @fold_assume_alignment_chain(%0: memref<128xf32>) -> memref<128xf32> {
+ // CHECK: %[[ALIGN:.+]] = memref.assume_alignment %[[ARG0]], 16
+ %1 = memref.assume_alignment %0, 16 : memref<128xf32>
+ // CHECK-NOT: memref.assume_alignment
+ %2 = memref.assume_alignment %1, 16 : memref<128xf32>
+ // CHECK: return %[[ALIGN]]
+ return %2 : memref<128xf32>
+}
More information about the Mlir-commits
mailing list