[Mlir-commits] [mlir] [MLIR] Allow `constFoldBinaryOp` to fold `(T1, T1) -> T2` (PR #151410)
Matthias Guenther
llvmlistbot at llvm.org
Tue Aug 5 15:54:46 PDT 2025
https://github.com/mrguenther updated https://github.com/llvm/llvm-project/pull/151410
>From 4b08b0b35049247060ea2bfbdbfc159cc2041cab Mon Sep 17 00:00:00 2001
From: mrguenther <mrguenther at google.com>
Date: Wed, 30 Jul 2025 14:54:24 -0700
Subject: [PATCH 1/3] Allow `constFoldBinaryOp` to fold `(T1, T1) -> T2`
The `constFoldBinaryOp` helper function had limited support for
different input and output types, but the static type of the underlying
value (e.g. `APInt`) had to match between the inputs and the output.
This worked fine for int comparisons of the form `(intN, intN) -> int1`,
as the static type signature was `(APInt, APInt) -> APInt`. However,
float comparisons map `(floatN, floatN) -> int1`, with a static type
signature of `(APFloat, APFloat) -> APInt`. This use case wasn't
supported by `constFoldBinaryOp`.
`constFoldBinaryOp` now accepts an optional template argument overriding
the return type in case it differs from the input type. If the new
template argument isn't provided, the default behavior is unchanged
(i.e. the return type will be assumed to match the input type).
---
mlir/include/mlir/Dialect/CommonFolders.h | 53 +++++++++++------
mlir/unittests/Dialect/CMakeLists.txt | 1 +
mlir/unittests/Dialect/CommonFoldersTest.cpp | 61 ++++++++++++++++++++
3 files changed, 96 insertions(+), 19 deletions(-)
create mode 100644 mlir/unittests/Dialect/CommonFoldersTest.cpp
diff --git a/mlir/include/mlir/Dialect/CommonFolders.h b/mlir/include/mlir/Dialect/CommonFolders.h
index b5a12426aff80..79832139b8610 100644
--- a/mlir/include/mlir/Dialect/CommonFolders.h
+++ b/mlir/include/mlir/Dialect/CommonFolders.h
@@ -15,10 +15,16 @@
#ifndef MLIR_DIALECT_COMMONFOLDERS_H
#define MLIR_DIALECT_COMMONFOLDERS_H
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Types.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
+
+#include <cassert>
+#include <cstddef>
#include <optional>
namespace mlir {
@@ -30,11 +36,13 @@ class PoisonAttr;
/// Uses `resultType` for the type of the returned attribute.
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
/// which will be directly propagated to result.
-template <class AttrElementT,
+template <class AttrElementT, //
class ElementValueT = typename AttrElementT::ValueType,
class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
class CalculationT = function_ref<
- std::optional<ElementValueT>(ElementValueT, ElementValueT)>>
+ std::optional<ResultElementValueT>(ElementValueT, ElementValueT)>>
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
Type resultType,
CalculationT &&calculate) {
@@ -65,7 +73,7 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
if (!calRes)
return {};
- return AttrElementT::get(resultType, *calRes);
+ return ResultAttrElementT::get(resultType, *calRes);
}
if (isa<SplatElementsAttr>(operands[0]) &&
@@ -99,7 +107,7 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
return {};
auto lhsIt = *maybeLhsIt;
auto rhsIt = *maybeRhsIt;
- SmallVector<ElementValueT, 4> elementResults;
+ SmallVector<ResultElementValueT, 4> elementResults;
elementResults.reserve(lhs.getNumElements());
for (size_t i = 0, e = lhs.getNumElements(); i < e; ++i, ++lhsIt, ++rhsIt) {
auto elementResult = calculate(*lhsIt, *rhsIt);
@@ -119,11 +127,13 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
/// attribute.
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
/// which will be directly propagated to result.
-template <class AttrElementT,
+template <class AttrElementT, //
class ElementValueT = typename AttrElementT::ValueType,
class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
class CalculationT = function_ref<
- std::optional<ElementValueT>(ElementValueT, ElementValueT)>>
+ std::optional<ResultElementValueT>(ElementValueT, ElementValueT)>>
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
CalculationT &&calculate) {
assert(operands.size() == 2 && "binary op takes two operands");
@@ -153,36 +163,41 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
return {};
return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT, ResultElementValueT,
CalculationT>(
operands, lhsType, std::forward<CalculationT>(calculate));
}
template <class AttrElementT,
class ElementValueT = typename AttrElementT::ValueType,
- class PoisonAttr = void,
+ class PoisonAttr = void, //
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
class CalculationT =
- function_ref<ElementValueT(ElementValueT, ElementValueT)>>
+ function_ref<ResultElementValueT(ElementValueT, ElementValueT)>>
Attribute constFoldBinaryOp(ArrayRef<Attribute> operands, Type resultType,
CalculationT &&calculate) {
- return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
+ return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT>(
operands, resultType,
- [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> {
- return calculate(a, b);
- });
+ [&](ElementValueT a, ElementValueT b)
+ -> std::optional<ResultElementValueT> { return calculate(a, b); });
}
-template <class AttrElementT,
+template <class AttrElementT, //
class ElementValueT = typename AttrElementT::ValueType,
class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
class CalculationT =
- function_ref<ElementValueT(ElementValueT, ElementValueT)>>
+ function_ref<ResultElementValueT(ElementValueT, ElementValueT)>>
Attribute constFoldBinaryOp(ArrayRef<Attribute> operands,
CalculationT &&calculate) {
- return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
+ return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT>(
operands,
- [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> {
- return calculate(a, b);
- });
+ [&](ElementValueT a, ElementValueT b)
+ -> std::optional<ResultElementValueT> { return calculate(a, b); });
}
/// Performs constant folding `calculate` with element-wise behavior on the one
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index aea247547473d..bbe1c021ec4bf 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -1,5 +1,6 @@
add_mlir_unittest(MLIRDialectTests
BroadcastShapeTest.cpp
+ CommonFoldersTest.cpp
)
mlir_target_link_libraries(MLIRDialectTests
PRIVATE
diff --git a/mlir/unittests/Dialect/CommonFoldersTest.cpp b/mlir/unittests/Dialect/CommonFoldersTest.cpp
new file mode 100644
index 0000000000000..31d9c592dabc1
--- /dev/null
+++ b/mlir/unittests/Dialect/CommonFoldersTest.cpp
@@ -0,0 +1,61 @@
+//===- CommonFoldersTest.cpp - tests for folder-pattern helper templates --===//
+//
+// 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/Dialect/CommonFolders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/Casting.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace mlir {
+namespace {
+
+using ::llvm::APFloat;
+using ::llvm::APInt;
+using ::mlir::constFoldBinaryOp;
+using ::mlir::DenseElementsAttr;
+using ::mlir::Float32Type;
+using ::mlir::FloatAttr;
+using ::mlir::IntegerAttr;
+using ::mlir::IntegerType;
+using ::mlir::MLIRContext;
+using ::mlir::RankedTensorType;
+using ::testing::ElementsAre;
+
+APInt floatLessThan(APFloat lhs, APFloat rhs) { return APInt(1, lhs < rhs); }
+
+TEST(CommonFoldersTest, FoldFloatComparisonToBoolean) {
+ MLIRContext context;
+ auto vector4xf32 = RankedTensorType::get({4}, Float32Type::get(&context));
+
+ auto lhs = DenseElementsAttr::get(vector4xf32, {-12.9f, 0.0f, 42.5f, -0.01f});
+ auto rhs = DenseElementsAttr::get(vector4xf32, {0.0f, 0.0f, 0.0f, 0.0f});
+
+ auto result = llvm::dyn_cast<DenseElementsAttr>(
+ constFoldBinaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(
+ {lhs, rhs}, RankedTensorType::get({4}, IntegerType::get(&context, 1)),
+ floatLessThan));
+ ASSERT_TRUE(result);
+
+ auto resultElementType = result.getElementType();
+ EXPECT_TRUE(resultElementType.isInteger(1));
+
+ const APInt i1True = APInt(1, true);
+ const APInt i1False = APInt(1, false);
+
+ EXPECT_THAT(result.getValues<APInt>(),
+ ElementsAre(i1True, i1False, i1False, i1True));
+}
+
+} // namespace
+} // namespace mlir
>From 8deccc197774c11793c58524bd86cedd8139260f Mon Sep 17 00:00:00 2001
From: mrguenther <mrguenther at google.com>
Date: Tue, 5 Aug 2025 15:37:15 -0700
Subject: [PATCH 2/3] Support unary ops; test with LIT instead of C++
---
mlir/include/mlir/Dialect/CommonFolders.h | 94 +++++++++++++++++---
mlir/test/Dialect/common_folders.mlir | 22 +++++
mlir/test/lib/Dialect/Test/TestOps.td | 20 +++++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 81 +++++++++++++++++
mlir/unittests/Dialect/CMakeLists.txt | 1 -
mlir/unittests/Dialect/CommonFoldersTest.cpp | 61 -------------
6 files changed, 203 insertions(+), 76 deletions(-)
create mode 100644 mlir/test/Dialect/common_folders.mlir
delete mode 100644 mlir/unittests/Dialect/CommonFoldersTest.cpp
diff --git a/mlir/include/mlir/Dialect/CommonFolders.h b/mlir/include/mlir/Dialect/CommonFolders.h
index 79832139b8610..113765157946d 100644
--- a/mlir/include/mlir/Dialect/CommonFolders.h
+++ b/mlir/include/mlir/Dialect/CommonFolders.h
@@ -149,14 +149,14 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
return operands[1];
}
- auto getResultType = [](Attribute attr) -> Type {
+ auto getAttrType = [](Attribute attr) -> Type {
if (auto typed = dyn_cast_or_null<TypedAttr>(attr))
return typed.getType();
return {};
};
- Type lhsType = getResultType(operands[0]);
- Type rhsType = getResultType(operands[1]);
+ Type lhsType = getAttrType(operands[0]);
+ Type rhsType = getAttrType(operands[1]);
if (!lhsType || !rhsType)
return {};
if (lhsType != rhsType)
@@ -202,16 +202,20 @@ Attribute constFoldBinaryOp(ArrayRef<Attribute> operands,
/// Performs constant folding `calculate` with element-wise behavior on the one
/// attributes in `operands` and returns the result if possible.
+/// Uses `resultType` for the type of the returned attribute.
/// Optional PoisonAttr template argument allows to specify 'poison' attribute
/// which will be directly propagated to result.
-template <class AttrElementT,
+template <class AttrElementT, //
class ElementValueT = typename AttrElementT::ValueType,
class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
class CalculationT =
- function_ref<std::optional<ElementValueT>(ElementValueT)>>
+ function_ref<std::optional<ResultElementValueT>(ElementValueT)>>
Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
+ Type resultType,
CalculationT &&calculate) {
- if (!llvm::getSingleElement(operands))
+ if (!resultType || !llvm::getSingleElement(operands))
return {};
static_assert(
@@ -229,7 +233,7 @@ Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
auto res = calculate(op.getValue());
if (!res)
return {};
- return AttrElementT::get(op.getType(), *res);
+ return ResultAttrElementT::get(resultType, *res);
}
if (isa<SplatElementsAttr>(operands[0])) {
// Both operands are splats so we can avoid expanding the values out and
@@ -239,7 +243,7 @@ Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
auto elementResult = calculate(op.getSplatValue<ElementValueT>());
if (!elementResult)
return {};
- return DenseElementsAttr::get(op.getType(), *elementResult);
+ return DenseElementsAttr::get(cast<ShapedType>(resultType), *elementResult);
} else if (isa<ElementsAttr>(operands[0])) {
// Operands are ElementsAttr-derived; perform an element-wise fold by
// expanding the values.
@@ -249,7 +253,7 @@ Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
if (!maybeOpIt)
return {};
auto opIt = *maybeOpIt;
- SmallVector<ElementValueT> elementResults;
+ SmallVector<ResultElementValueT> elementResults;
elementResults.reserve(op.getNumElements());
for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
auto elementResult = calculate(*opIt);
@@ -257,19 +261,81 @@ Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
return {};
elementResults.push_back(*elementResult);
}
- return DenseElementsAttr::get(op.getShapedType(), elementResults);
+ return DenseElementsAttr::get(cast<ShapedType>(resultType), elementResults);
}
return {};
}
-template <class AttrElementT,
+/// Performs constant folding `calculate` with element-wise behavior on the one
+/// attributes in `operands` and returns the result if possible.
+/// Uses the operand element type for the element type of the returned
+/// attribute.
+/// Optional PoisonAttr template argument allows to specify 'poison' attribute
+/// which will be directly propagated to result.
+template <class AttrElementT, //
+ class ElementValueT = typename AttrElementT::ValueType,
+ class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
+ class CalculationT =
+ function_ref<std::optional<ResultElementValueT>(ElementValueT)>>
+Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
+ CalculationT &&calculate) {
+ if (!llvm::getSingleElement(operands))
+ return {};
+
+ static_assert(
+ std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
+ "PoisonAttr is undefined, either add a dependency on UB dialect or pass "
+ "void as template argument to opt-out from poison semantics.");
+ if constexpr (!std::is_void_v<PoisonAttr>) {
+ if (isa<PoisonAttr>(operands[0]))
+ return operands[0];
+ }
+
+ auto getAttrType = [](Attribute attr) -> Type {
+ if (auto typed = dyn_cast_or_null<TypedAttr>(attr))
+ return typed.getType();
+ return {};
+ };
+
+ Type operandType = getAttrType(operands[0]);
+ if (!operandType)
+ return {};
+
+ return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT, ResultElementValueT,
+ CalculationT>(
+ operands, operandType, std::forward<CalculationT>(calculate));
+}
+
+template <class AttrElementT, //
+ class ElementValueT = typename AttrElementT::ValueType,
+ class PoisonAttr = ub::PoisonAttr,
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
+ class CalculationT = function_ref<ResultElementValueT(ElementValueT)>>
+Attribute constFoldUnaryOp(ArrayRef<Attribute> operands, Type resultType,
+ CalculationT &&calculate) {
+ return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT>(
+ operands, resultType,
+ [&](ElementValueT a) -> std::optional<ResultElementValueT> {
+ return calculate(a);
+ });
+}
+
+template <class AttrElementT, //
class ElementValueT = typename AttrElementT::ValueType,
class PoisonAttr = ub::PoisonAttr,
- class CalculationT = function_ref<ElementValueT(ElementValueT)>>
+ class ResultAttrElementT = AttrElementT,
+ class ResultElementValueT = typename ResultAttrElementT::ValueType,
+ class CalculationT = function_ref<ResultElementValueT(ElementValueT)>>
Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,
CalculationT &&calculate) {
- return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
- operands, [&](ElementValueT a) -> std::optional<ElementValueT> {
+ return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+ ResultAttrElementT>(
+ operands, [&](ElementValueT a) -> std::optional<ResultElementValueT> {
return calculate(a);
});
}
diff --git a/mlir/test/Dialect/common_folders.mlir b/mlir/test/Dialect/common_folders.mlir
new file mode 100644
index 0000000000000..92598b4937552
--- /dev/null
+++ b/mlir/test/Dialect/common_folders.mlir
@@ -0,0 +1,22 @@
+// RUN: mlir-opt %s --test-fold-type-converting-op --split-input-file | FileCheck %s
+
+// CHECK-LABEL: @test_fold_unary_op_f32_to_si32(
+func.func @test_fold_unary_op_f32_to_si32() -> tensor<4x2xsi32> {
+ // CHECK-NEXT: %[[POSITIVE_ONE:.*]] = arith.constant dense<1> : tensor<4x2xsi32>
+ // CHECK-NEXT: return %[[POSITIVE_ONE]] : tensor<4x2xsi32>
+ %operand = arith.constant dense<5.1> : tensor<4x2xf32>
+ %sign = test.sign %operand : (tensor<4x2xf32>) -> tensor<4x2xsi32>
+ return %sign : tensor<4x2xsi32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_fold_binary_op_f32_to_i1(
+func.func @test_fold_binary_op_f32_to_i1() -> tensor<8xi1> {
+ // CHECK-NEXT: %[[FALSE:.*]] = arith.constant dense<false> : tensor<8xi1>
+ // CHECK-NEXT: return %[[FALSE]] : tensor<8xi1>
+ %lhs = arith.constant dense<5.1> : tensor<8xf32>
+ %rhs = arith.constant dense<4.2> : tensor<8xf32>
+ %less_than = test.less_than %lhs, %rhs : (tensor<8xf32>, tensor<8xf32>) -> tensor<8xi1>
+ return %less_than : tensor<8xi1>
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 2eaad552a7a3a..2dd8b61e410d0 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1169,6 +1169,26 @@ def OpP : TEST_Op<"op_p"> {
let results = (outs I32);
}
+// Test constant-folding a pattern that maps (F32, F32) -> I32.
+def SignOp : TEST_Op<"sign", [SameOperandsAndResultShape]> {
+ let arguments = (ins RankedTensorOf<[F32]>:$operand);
+ let results = (outs RankedTensorOf<[SI32]>:$result);
+
+ let assemblyFormat = [{
+ $operand attr-dict `:` functional-type(operands, results)
+ }];
+}
+
+// Test constant-folding a pattern that maps (F32, F32) -> I32.
+def LessThanOp : TEST_Op<"less_than", [SameOperandsAndResultShape]> {
+ let arguments = (ins RankedTensorOf<[F32]>:$lhs, RankedTensorOf<[F32]>:$rhs);
+ let results = (outs RankedTensorOf<[I1]>:$result);
+
+ let assemblyFormat = [{
+ $lhs `,` $rhs attr-dict `:` functional-type(operands, results)
+ }];
+}
+
// Test same operand name enforces equality condition check.
def TestEqualArgsPattern : Pat<(OpN $a, $a), (OpO $a)>;
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index eda618f5b09c6..1ad1b6163374d 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -10,6 +10,7 @@
#include "TestOps.h"
#include "TestTypes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/CommonFolders.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -202,6 +203,66 @@ struct HoistEligibleOps : public OpRewritePattern<test::OneRegionOp> {
}
};
+struct FoldSignOpF32ToSI32 : public OpRewritePattern<test::SignOp> {
+ using OpRewritePattern<test::SignOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(test::SignOp op,
+ PatternRewriter &rewriter) const override {
+ if (op->getNumOperands() != 1 || op->getNumResults() != 1)
+ return failure();
+
+ TypedAttr operandAttr;
+ matchPattern(op->getOperand(0), m_Constant(&operandAttr));
+ if (!operandAttr)
+ return failure();
+
+ TypedAttr res = cast_or_null<TypedAttr>(
+ constFoldUnaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(
+ operandAttr, op.getType(), [](APFloat operand) -> APSInt {
+ static const APFloat zero(0.0f);
+ int operandSign = 0;
+ if (operand != zero)
+ operandSign = (operand < zero) ? -1 : +1;
+ return APSInt(APInt(32, operandSign), false);
+ }));
+ if (!res)
+ return failure();
+
+ rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, res);
+ return success();
+ }
+};
+
+struct FoldLessThanOpF32ToI1 : public OpRewritePattern<test::LessThanOp> {
+ using OpRewritePattern<test::LessThanOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(test::LessThanOp op,
+ PatternRewriter &rewriter) const override {
+ if (op->getNumOperands() != 2 || op->getNumResults() != 1)
+ return failure();
+
+ TypedAttr lhsAttr;
+ TypedAttr rhsAttr;
+ matchPattern(op->getOperand(0), m_Constant(&lhsAttr));
+ matchPattern(op->getOperand(1), m_Constant(&rhsAttr));
+
+ if (!lhsAttr || !rhsAttr)
+ return failure();
+
+ Attribute operandAttrs[2] = {lhsAttr, rhsAttr};
+ TypedAttr res = cast_or_null<TypedAttr>(
+ constFoldBinaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(
+ operandAttrs, op.getType(), [](APFloat lhs, APFloat rhs) -> APInt {
+ return APInt(1, lhs < rhs);
+ }));
+ if (!res)
+ return failure();
+
+ rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, res);
+ return success();
+ }
+};
+
/// This pattern moves "test.move_before_parent_op" before the parent op.
struct MoveBeforeParentOp : public RewritePattern {
MoveBeforeParentOp(MLIRContext *context)
@@ -2181,6 +2242,24 @@ struct TestSelectiveReplacementPatternDriver
(void)applyPatternsGreedily(getOperation(), std::move(patterns));
}
};
+
+struct TestFoldTypeConvertingOp
+ : public PassWrapper<TestFoldTypeConvertingOp, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestFoldTypeConvertingOp)
+
+ StringRef getArgument() const final { return "test-fold-type-converting-op"; }
+ StringRef getDescription() const final {
+ return "Test helper functions for folding ops whose input and output types "
+ "differ, e.g. float comparisons of the form `(f32, f32) -> i1`.";
+ }
+ void runOnOperation() override {
+ MLIRContext *context = &getContext();
+ mlir::RewritePatternSet patterns(context);
+ patterns.add<FoldSignOpF32ToSI32, FoldLessThanOpF32ToI1>(context);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ signalPassFailure();
+ }
+};
} // namespace
//===----------------------------------------------------------------------===//
@@ -2211,6 +2290,8 @@ void registerPatternsTestPass() {
PassRegistration<TestMergeBlocksPatternDriver>();
PassRegistration<TestSelectiveReplacementPatternDriver>();
+
+ PassRegistration<TestFoldTypeConvertingOp>();
}
} // namespace test
} // namespace mlir
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index bbe1c021ec4bf..aea247547473d 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -1,6 +1,5 @@
add_mlir_unittest(MLIRDialectTests
BroadcastShapeTest.cpp
- CommonFoldersTest.cpp
)
mlir_target_link_libraries(MLIRDialectTests
PRIVATE
diff --git a/mlir/unittests/Dialect/CommonFoldersTest.cpp b/mlir/unittests/Dialect/CommonFoldersTest.cpp
deleted file mode 100644
index 31d9c592dabc1..0000000000000
--- a/mlir/unittests/Dialect/CommonFoldersTest.cpp
+++ /dev/null
@@ -1,61 +0,0 @@
-//===- CommonFoldersTest.cpp - tests for folder-pattern helper templates --===//
-//
-// 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/Dialect/CommonFolders.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/MLIRContext.h"
-#include "llvm/ADT/APFloat.h"
-#include "llvm/ADT/APInt.h"
-#include "llvm/Support/Casting.h"
-#include "gmock/gmock.h"
-#include "gtest/gtest.h"
-
-namespace mlir {
-namespace {
-
-using ::llvm::APFloat;
-using ::llvm::APInt;
-using ::mlir::constFoldBinaryOp;
-using ::mlir::DenseElementsAttr;
-using ::mlir::Float32Type;
-using ::mlir::FloatAttr;
-using ::mlir::IntegerAttr;
-using ::mlir::IntegerType;
-using ::mlir::MLIRContext;
-using ::mlir::RankedTensorType;
-using ::testing::ElementsAre;
-
-APInt floatLessThan(APFloat lhs, APFloat rhs) { return APInt(1, lhs < rhs); }
-
-TEST(CommonFoldersTest, FoldFloatComparisonToBoolean) {
- MLIRContext context;
- auto vector4xf32 = RankedTensorType::get({4}, Float32Type::get(&context));
-
- auto lhs = DenseElementsAttr::get(vector4xf32, {-12.9f, 0.0f, 42.5f, -0.01f});
- auto rhs = DenseElementsAttr::get(vector4xf32, {0.0f, 0.0f, 0.0f, 0.0f});
-
- auto result = llvm::dyn_cast<DenseElementsAttr>(
- constFoldBinaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(
- {lhs, rhs}, RankedTensorType::get({4}, IntegerType::get(&context, 1)),
- floatLessThan));
- ASSERT_TRUE(result);
-
- auto resultElementType = result.getElementType();
- EXPECT_TRUE(resultElementType.isInteger(1));
-
- const APInt i1True = APInt(1, true);
- const APInt i1False = APInt(1, false);
-
- EXPECT_THAT(result.getValues<APInt>(),
- ElementsAre(i1True, i1False, i1False, i1True));
-}
-
-} // namespace
-} // namespace mlir
>From 9a734d009ef7c3aad3586cc98b1b61c4371fc34c Mon Sep 17 00:00:00 2001
From: mrguenther <mrguenther at google.com>
Date: Tue, 5 Aug 2025 15:53:11 -0700
Subject: [PATCH 3/3] Fix a couple of typos
---
mlir/test/lib/Dialect/Test/TestOps.td | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 2dd8b61e410d0..f3554e8579d29 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1169,7 +1169,7 @@ def OpP : TEST_Op<"op_p"> {
let results = (outs I32);
}
-// Test constant-folding a pattern that maps (F32, F32) -> I32.
+// Test constant-folding a pattern that maps `(F32) -> SI32`.
def SignOp : TEST_Op<"sign", [SameOperandsAndResultShape]> {
let arguments = (ins RankedTensorOf<[F32]>:$operand);
let results = (outs RankedTensorOf<[SI32]>:$result);
@@ -1179,7 +1179,7 @@ def SignOp : TEST_Op<"sign", [SameOperandsAndResultShape]> {
}];
}
-// Test constant-folding a pattern that maps (F32, F32) -> I32.
+// Test constant-folding a pattern that maps `(F32, F32) -> I1`.
def LessThanOp : TEST_Op<"less_than", [SameOperandsAndResultShape]> {
let arguments = (ins RankedTensorOf<[F32]>:$lhs, RankedTensorOf<[F32]>:$rhs);
let results = (outs RankedTensorOf<[I1]>:$result);
More information about the Mlir-commits
mailing list