[Mlir-commits] [mlir] [mlir][Math] Add pass to legalize math functions to f32-or-higher (PR #78361)

[Krzysztof Drewniak]

================
@@ -0,0 +1,118 @@
+//===- LegalizeToF32.cpp - Legalize functions on small floats ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements legalizing math operations on small floating-point
+// types through arith.extf and arith.truncf.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace mlir::math {
+#define GEN_PASS_DEF_MATHLEGALIZETOF32
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+} // namespace mlir::math
+
+using namespace mlir;
+namespace {
+struct LegalizeToF32RewritePattern final : ConversionPattern {
+  LegalizeToF32RewritePattern(TypeConverter &converter, MLIRContext *context)
+      : ConversionPattern(converter, MatchAnyOpTypeTag{}, 1, context) {}
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+struct LegalizeToF32Pass final
+    : mlir::math::impl::MathLegalizeToF32Base<LegalizeToF32Pass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void mlir::math::populateLegalizeToF32TypeConverter(
+    TypeConverter &typeConverter) {
+  typeConverter.addConversion(
+      [](Type type) -> std::optional<Type> { return type; });
+  typeConverter.addConversion([](FloatType type) -> std::optional<Type> {
+    if (type.getWidth() < 32)
+      return Float32Type::get(type.getContext());
+    return std::nullopt;
+  });
+  typeConverter.addConversion([](ShapedType type) -> std::optional<Type> {
+    if (auto elemTy = dyn_cast<FloatType>(type.getElementType()))
+      return type.clone(Float32Type::get(type.getContext()));
+    return std::nullopt;
+  });
+  typeConverter.addTargetMaterialization(
+      [](OpBuilder &b, Type target, ValueRange input, Location loc) {
+        return b.create<arith::ExtFOp>(loc, target, input);
+      });
+}
+
+void mlir::math::populateLegalizeToF32ConversionTarget(
+    ConversionTarget &target, TypeConverter &typeConverter) {
+  target.addDynamicallyLegalDialect<MathDialect>(
+      [&typeConverter](Operation *op) -> bool {
+        return typeConverter.isLegal(op);
+      });
+  target.addLegalOp<FmaOp>();
+  target.addLegalOp<arith::ExtFOp, arith::TruncFOp>();
+}
+
+LogicalResult LegalizeToF32RewritePattern::matchAndRewrite(
+    Operation *op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  Location loc = op->getLoc();
+  const TypeConverter *converter = getTypeConverter();
+  if (converter->isLegal(op))
+    return rewriter.notifyMatchFailure(loc, "op already legal");
+  OperationState newOp(loc, op->getName());
+  newOp.addOperands(operands);
+
+  SmallVector<Type> newResultTypes;
+  if (failed(converter->convertTypes(op->getResultTypes(), newResultTypes)))
+    return rewriter.notifyMatchFailure(loc, "couldn't convert return types");
+  newOp.addTypes(newResultTypes);
+  newOp.addAttributes(op->getAttrs());
+  Operation *legalized = rewriter.create(newOp);
----------------
krzysz00 wrote:

I mean, I had the `OperationState` version so I could copy over attributes, properties, etc. while creating a new operation.

I'm not sure `clone()` + mutations really buys us anything over just the mutations. And looking at dialect conversion's resets, the only rollback you'd want to be able to do is `cancelOpModification` which will reset the operands (though not result types, apparently)

https://github.com/llvm/llvm-project/pull/78361