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).

Author: mrguenther

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

mlir/unittests/Dialect/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
+  CommonFoldersTest.cpp
 )
 mlir_target_link_libraries(MLIRDialectTests
   PRIVATE

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