make everyone happy

makslevental · makslevental · commit c79a9ef43814 · 2025-11-10T11:32:31.000-08:00
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
@@ -199,7 +199,6 @@ def ArithToAPFloatConversionPass
     floating-point arithmetic operations.
   }];
   let dependentDialects = ["func::FuncDialect"];
-  let options = [];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Func/Utils/Utils.h b/mlir/include/mlir/Dialect/Func/Utils/Utils.h
@@ -60,13 +60,12 @@ mlir::FailureOr<std::pair<mlir::func::FuncOp, mlir::func::CallOp>>
 deduplicateArgsOfFuncOp(mlir::RewriterBase &rewriter, mlir::func::FuncOp funcOp,
                         mlir::ModuleOp moduleOp);
 
-/// Create a FuncOp with signature `resultTypes`(`paramTypes`)` and name `name`.
-/// Return a failure if the FuncOp found has unexpected signature.
-FailureOr<FuncOp>
-lookupOrCreateFnDecl(OpBuilder &b, Operation *moduleOp, StringRef name,
-                     ArrayRef<Type> paramTypes = {},
-                     ArrayRef<Type> resultTypes = {}, bool setPrivate = false,
-                     SymbolTableCollection *symbolTables = nullptr);
+/// Look up a FuncOp with signature `resultTypes`(`paramTypes`)` and name
+/// `name`. Return a failure if the FuncOp is found but with a different
+/// signature.
+FailureOr<FuncOp> lookupFnDecl(SymbolOpInterface symTable, StringRef name,
+                               FunctionType funcT,
+                               SymbolTableCollection *symbolTables = nullptr);
 
 } // namespace func
 } // namespace mlir
diff --git a/mlir/lib/Conversion/ArithToAPFloat/ArithToAPFloat.cpp b/mlir/lib/Conversion/ArithToAPFloat/ArithToAPFloat.cpp
@@ -15,6 +15,9 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "arith-to-apfloat"
 
 namespace mlir {
 #define GEN_PASS_DEF_ARITHTOAPFLOATCONVERSIONPASS
@@ -24,6 +27,22 @@ namespace mlir {
 using namespace mlir;
 using namespace mlir::func;
 
+static FuncOp createFnDecl(OpBuilder &b, SymbolOpInterface symTable,
+                           StringRef name, FunctionType funcT, bool setPrivate,
+                           SymbolTableCollection *symbolTables = nullptr) {
+  OpBuilder::InsertionGuard g(b);
+  assert(!symTable->getRegion(0).empty() && "expected non-empty region");
+  b.setInsertionPointToStart(&symTable->getRegion(0).front());
+  FuncOp funcOp = FuncOp::create(b, symTable->getLoc(), name, funcT);
+  if (setPrivate)
+    funcOp.setPrivate();
+  if (symbolTables) {
+    SymbolTable &symbolTable = symbolTables->getSymbolTable(symTable);
+    symbolTable.insert(funcOp, symTable->getRegion(0).front().begin());
+  }
+  return funcOp;
+}
+
 /// Helper function to look up or create the symbol for a runtime library
 /// function for a binary arithmetic operation.
 ///
@@ -34,34 +53,42 @@ using namespace mlir::func;
 /// This function will return a failure if the function is found but has an
 /// unexpected signature.
 ///
-static FailureOr<Operation *>
-lookupOrCreateBinaryFn(OpBuilder &b, Operation *moduleOp, StringRef name,
+static FailureOr<FuncOp>
+lookupOrCreateBinaryFn(OpBuilder &b, SymbolOpInterface symTable, StringRef name,
                        SymbolTableCollection *symbolTables = nullptr) {
-  auto i32Type = IntegerType::get(moduleOp->getContext(), 32);
-  auto i64Type = IntegerType::get(moduleOp->getContext(), 64);
-  return lookupOrCreateFnDecl(b, moduleOp,
-                              (llvm::Twine("_mlir_apfloat_") + name).str(),
-                              {i32Type, i64Type, i64Type}, {i64Type},
-                              /*setPrivate=*/true, symbolTables);
+  auto i32Type = IntegerType::get(symTable->getContext(), 32);
+  auto i64Type = IntegerType::get(symTable->getContext(), 64);
+
+  std::string funcName = (llvm::Twine("__mlir_apfloat_") + name).str();
+  FunctionType funcT =
+      FunctionType::get(b.getContext(), {i32Type, i64Type, i64Type}, {i64Type});
+  FailureOr<FuncOp> func =
+      lookupFnDecl(symTable, funcName, funcT, symbolTables);
+  // Failed due to type mismatch.
+  if (failed(func))
+    return func;
+  // Successfully matched existing decl.
+  if (*func)
+    return *func;
+
+  return createFnDecl(b, symTable, funcName, funcT,
+                      /*setPrivate=*/true, symbolTables);
 }
 
 /// Rewrite a binary arithmetic operation to an APFloat function call.
 template <typename OpTy, const char *APFloatName>
-struct ArithOpToAPFloatConversion final : OpRewritePattern<OpTy> {
-  using OpRewritePattern<OpTy>::OpRewritePattern;
+struct BinaryArithOpToAPFloatConversion final : OpRewritePattern<OpTy> {
+  BinaryArithOpToAPFloatConversion(MLIRContext *context, PatternBenefit benefit,
+                                   SymbolOpInterface symTable)
+      : OpRewritePattern<OpTy>(context, benefit), symTable(symTable) {};
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
-    auto moduleOp = op->template getParentOfType<ModuleOp>();
-    if (!moduleOp) {
-      return rewriter.notifyMatchFailure(
-          op, "arith op must be contained within a builtin.module");
-    }
     // Get APFloat function from runtime library.
-    FailureOr<Operation *> fn =
-        lookupOrCreateBinaryFn(rewriter, moduleOp, APFloatName);
+    FailureOr<FuncOp> fn =
+        lookupOrCreateBinaryFn(rewriter, symTable, APFloatName);
     if (failed(fn))
-      return op->emitError("failed to lookup or create APFloat function");
+      return fn;
 
     rewriter.setInsertionPoint(op);
     // Cast operands to 64-bit integers.
@@ -94,6 +121,8 @@ struct ArithOpToAPFloatConversion final : OpRewritePattern<OpTy> {
         op, arith::BitcastOp::create(rewriter, loc, floatTy, truncatedBits));
     return success();
   }
+
+  SymbolOpInterface symTable;
 };
 
 namespace {
@@ -109,12 +138,24 @@ struct ArithToAPFloatConversionPass final
     static const char multiply[] = "multiply";
     static const char divide[] = "divide";
     static const char remainder[] = "remainder";
-    patterns.add<ArithOpToAPFloatConversion<arith::AddFOp, add>,
-                 ArithOpToAPFloatConversion<arith::SubFOp, subtract>,
-                 ArithOpToAPFloatConversion<arith::MulFOp, multiply>,
-                 ArithOpToAPFloatConversion<arith::DivFOp, divide>,
-                 ArithOpToAPFloatConversion<arith::RemFOp, remainder>>(context);
+    patterns.add<BinaryArithOpToAPFloatConversion<arith::AddFOp, add>,
+                 BinaryArithOpToAPFloatConversion<arith::SubFOp, subtract>,
+                 BinaryArithOpToAPFloatConversion<arith::MulFOp, multiply>,
+                 BinaryArithOpToAPFloatConversion<arith::DivFOp, divide>,
+                 BinaryArithOpToAPFloatConversion<arith::RemFOp, remainder>>(
+        context, 1, getOperation());
+    LogicalResult result = success();
+    ScopedDiagnosticHandler scopedHandler(context, [&result](Diagnostic &diag) {
+      if (diag.getSeverity() == DiagnosticSeverity::Error) {
+        result = failure();
+      }
+      // NB: if you don't return failure, no other diag handlers will fire (see
+      // mlir/lib/IR/Diagnostics.cpp:DiagnosticEngineImpl::emit).
+      return failure();
+    });
     walkAndApplyPatterns(getOperation(), std::move(patterns));
+    if (failed(result))
+      return signalPassFailure();
   }
 };
 } // namespace
diff --git a/mlir/lib/Dialect/Func/Utils/Utils.cpp b/mlir/lib/Dialect/Func/Utils/Utils.cpp
@@ -256,44 +256,26 @@ func::deduplicateArgsOfFuncOp(RewriterBase &rewriter, func::FuncOp funcOp,
 }
 
 FailureOr<func::FuncOp>
-func::lookupOrCreateFnDecl(OpBuilder &b, Operation *moduleOp, StringRef name,
-                           ArrayRef<Type> paramTypes,
-                           ArrayRef<Type> resultTypes, bool setPrivate,
-                           SymbolTableCollection *symbolTables) {
-  assert(moduleOp->hasTrait<OpTrait::SymbolTable>() &&
-         "expected SymbolTable operation");
-
+func::lookupFnDecl(SymbolOpInterface symTable, StringRef name,
+                   FunctionType funcT, SymbolTableCollection *symbolTables) {
   FuncOp func;
   if (symbolTables) {
     func = symbolTables->lookupSymbolIn<FuncOp>(
-        moduleOp, StringAttr::get(moduleOp->getContext(), name));
+        symTable, StringAttr::get(symTable->getContext(), name));
   } else {
     func = llvm::dyn_cast_or_null<FuncOp>(
-        SymbolTable::lookupSymbolIn(moduleOp, name));
+        SymbolTable::lookupSymbolIn(symTable, name));
   }
 
-  FunctionType funcT =
-      FunctionType::get(b.getContext(), paramTypes, resultTypes);
-  // Assert the signature of the found function is same as expected
-  if (func) {
-    if (funcT != func.getFunctionType()) {
-      func.emitError("redefinition of function '")
-          << name << "' of different type " << funcT << " is prohibited";
-      return failure();
-    }
+  if (!func)
     return func;
-  }
 
-  OpBuilder::InsertionGuard g(b);
-  assert(!moduleOp->getRegion(0).empty() && "expected non-empty region");
-  b.setInsertionPointToStart(&moduleOp->getRegion(0).front());
-  FuncOp funcOp = FuncOp::create(b, moduleOp->getLoc(), name, funcT);
-  if (setPrivate)
-    funcOp.setPrivate();
-  if (symbolTables) {
-    SymbolTable &symbolTable = symbolTables->getSymbolTable(moduleOp);
-    symbolTable.insert(funcOp, moduleOp->getRegion(0).front().begin());
+  mlir::FunctionType foundFuncT = func.getFunctionType();
+  // Assert the signature of the found function is same as expected
+  if (funcT != foundFuncT) {
+    return func.emitError("matched function '")
+           << name << "' but with different type: " << foundFuncT
+           << " (expected " << funcT << ")";
   }
-
-  return funcOp;
+  return func;
 }
diff --git a/mlir/lib/ExecutionEngine/APFloatWrappers.cpp b/mlir/lib/ExecutionEngine/APFloatWrappers.cpp
@@ -29,7 +29,7 @@
 
 /// Binary operations without rounding mode.
 #define APFLOAT_BINARY_OP(OP)                                                  \
-  int64_t MLIR_APFLOAT_WRAPPERS_EXPORTED _mlir_apfloat_##OP(                   \
+  int64_t MLIR_APFLOAT_WRAPPERS_EXPORTED __mlir_apfloat_##OP(                   \
       int32_t semantics, uint64_t a, uint64_t b) {                             \
     const llvm::fltSemantics &sem = llvm::APFloatBase::EnumToSemantics(        \
         static_cast<llvm::APFloatBase::Semantics>(semantics));                 \
@@ -42,7 +42,7 @@
 
 /// Binary operations with rounding mode.
 #define APFLOAT_BINARY_OP_ROUNDING_MODE(OP, ROUNDING_MODE)                     \
-  int64_t MLIR_APFLOAT_WRAPPERS_EXPORTED _mlir_apfloat_##OP(                   \
+  int64_t MLIR_APFLOAT_WRAPPERS_EXPORTED __mlir_apfloat_##OP(                   \
       int32_t semantics, uint64_t a, uint64_t b) {                             \
     const llvm::fltSemantics &sem = llvm::APFloatBase::EnumToSemantics(        \
         static_cast<llvm::APFloatBase::Semantics>(semantics));                 \
diff --git a/mlir/test/Conversion/ArithToApfloat/arith-to-apfloat.mlir b/mlir/test/Conversion/ArithToApfloat/arith-to-apfloat.mlir
@@ -1,24 +1,45 @@
-// RUN: mlir-opt %s --convert-arith-to-apfloat -split-input-file | FileCheck %s
+// RUN: mlir-opt %s --convert-arith-to-apfloat -split-input-file -verify-diagnostics | FileCheck %s
 
-// CHECK-LABEL:   func.func private @_mlir_apfloat_add(i32, i64, i64) -> i64
+// CHECK-LABEL:   func.func private @__mlir_apfloat_add(i32, i64, i64) -> i64
 
 // CHECK-LABEL:   func.func @foo() -> f8E4M3FN {
 // CHECK:           %[[CONSTANT_0:.*]] = arith.constant 2.250000e+00 : f8E4M3FN
 // CHECK:           return %[[CONSTANT_0]] : f8E4M3FN
 // CHECK:         }
 
+// CHECK-LABEL:   func.func @bar() -> f6E3M2FN {
+// CHECK:           %[[CONSTANT_0:.*]] = arith.constant 3.000000e+00 : f6E3M2FN
+// CHECK:           return %[[CONSTANT_0]] : f6E3M2FN
+// CHECK:         }
+
+// Illustrate that both f8E4M3FN and f6E3M2FN calling the same __mlir_apfloat_add is fine
+// because each gets its own semantics enum and gets bitcast/extui/trunci to its own width.
 // CHECK-LABEL:   func.func @full_example() {
-// CHECK:           %[[cst:.*]] = arith.constant 1.375000e+00 : f8E4M3FN
-// CHECK:           %[[rhs:.*]] = call @foo() : () -> f8E4M3FN
-// CHECK:           %[[lhs_casted:.*]] = arith.bitcast %[[cst]] : f8E4M3FN to i8
-// CHECK:           %[[lhs_ext:.*]] = arith.extui %[[lhs_casted]] : i8 to i64
-// CHECK:           %[[rhs_casted:.*]] = arith.bitcast %[[rhs]] : f8E4M3FN to i8
-// CHECK:           %[[rhs_ext:.*]] = arith.extui %[[rhs_casted]] : i8 to i64
-// CHECK:           %[[c10_i32:.*]] = arith.constant 10 : i32
-// CHECK:           %[[res:.*]] = call @_mlir_apfloat_add(%[[c10_i32]], %[[lhs_ext]], %[[rhs_ext]]) : (i32, i64, i64) -> i64
-// CHECK:           %[[res_trunc:.*]] = arith.trunci %[[res]] : i64 to i8
-// CHECK:           %[[res_casted:.*]] = arith.bitcast %[[res_trunc]] : i8 to f8E4M3FN
-// CHECK:           vector.print %[[res_casted]] : f8E4M3FN
+// CHECK:           %[[CONSTANT_0:.*]] = arith.constant 1.375000e+00 : f8E4M3FN
+// CHECK:           %[[VAL_0:.*]] = call @foo() : () -> f8E4M3FN
+// CHECK:           %[[BITCAST_0:.*]] = arith.bitcast %[[CONSTANT_0]] : f8E4M3FN to i8
+// CHECK:           %[[EXTUI_0:.*]] = arith.extui %[[BITCAST_0]] : i8 to i64
+// CHECK:           %[[BITCAST_1:.*]] = arith.bitcast %[[VAL_0]] : f8E4M3FN to i8
+// CHECK:           %[[EXTUI_1:.*]] = arith.extui %[[BITCAST_1]] : i8 to i64
+//                  // fltSemantics semantics for f8E4M3FN
+// CHECK:           %[[CONSTANT_1:.*]] = arith.constant 10 : i32
+// CHECK:           %[[VAL_1:.*]] = call @__mlir_apfloat_add(%[[CONSTANT_1]], %[[EXTUI_0]], %[[EXTUI_1]]) : (i32, i64, i64) -> i64
+// CHECK:           %[[TRUNCI_0:.*]] = arith.trunci %[[VAL_1]] : i64 to i8
+// CHECK:           %[[BITCAST_2:.*]] = arith.bitcast %[[TRUNCI_0]] : i8 to f8E4M3FN
+// CHECK:           vector.print %[[BITCAST_2]] : f8E4M3FN
+
+// CHECK:           %[[CONSTANT_2:.*]] = arith.constant 2.500000e+00 : f6E3M2FN
+// CHECK:           %[[VAL_2:.*]] = call @bar() : () -> f6E3M2FN
+// CHECK:           %[[BITCAST_3:.*]] = arith.bitcast %[[CONSTANT_2]] : f6E3M2FN to i6
+// CHECK:           %[[EXTUI_2:.*]] = arith.extui %[[BITCAST_3]] : i6 to i64
+// CHECK:           %[[BITCAST_4:.*]] = arith.bitcast %[[VAL_2]] : f6E3M2FN to i6
+// CHECK:           %[[EXTUI_3:.*]] = arith.extui %[[BITCAST_4]] : i6 to i64
+//                  // fltSemantics semantics for f6E3M2FN
+// CHECK:           %[[CONSTANT_3:.*]] = arith.constant 16 : i32
+// CHECK:           %[[VAL_3:.*]] = call @__mlir_apfloat_add(%[[CONSTANT_3]], %[[EXTUI_2]], %[[EXTUI_3]]) : (i32, i64, i64) -> i64
+// CHECK:           %[[TRUNCI_1:.*]] = arith.trunci %[[VAL_3]] : i64 to i6
+// CHECK:           %[[BITCAST_5:.*]] = arith.bitcast %[[TRUNCI_1]] : i6 to f6E3M2FN
+// CHECK:           vector.print %[[BITCAST_5]] : f6E3M2FN
 // CHECK:           return
 // CHECK:         }
 
@@ -28,60 +49,79 @@ func.func @foo() -> f8E4M3FN {
   return %cst : f8E4M3FN
 }
 
+func.func @bar() -> f6E3M2FN {
+  %cst = arith.constant 3.2 : f6E3M2FN
+  return %cst : f6E3M2FN
+}
+
 func.func @full_example() {
   %a = arith.constant 1.4 : f8E4M3FN
   %b = func.call @foo() : () -> (f8E4M3FN)
   %c = arith.addf %a, %b : f8E4M3FN
-
   vector.print %c : f8E4M3FN
+
+  %d = arith.constant 2.4 : f6E3M2FN
+  %e = func.call @bar() : () -> (f6E3M2FN)
+  %f = arith.addf %d, %e : f6E3M2FN
+  vector.print %f : f6E3M2FN
   return
 }
 
 // -----
 
-// CHECK: func.func private @_mlir_apfloat_add(i32, i64, i64) -> i64
+// CHECK: func.func private @__mlir_apfloat_add(i32, i64, i64) -> i64
 // CHECK: %[[sem:.*]] = arith.constant 18 : i32
-// CHECK: call @_mlir_apfloat_add(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+// CHECK: call @__mlir_apfloat_add(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+func.func @addf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
+  %0 = arith.addf %arg0, %arg1 : f4E2M1FN
+  return
+}
+
+// -----
+
+// Test decl collision (different type)
+// expected-error@+1{{matched function '__mlir_apfloat_add' but with different type: '(i32, i32, f32) -> index' (expected '(i32, i64, i64) -> i64')}}
+func.func private @__mlir_apfloat_add(i32, i32, f32) -> index
 func.func @addf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
   %0 = arith.addf %arg0, %arg1 : f4E2M1FN
   return
 }
 
 // -----
 
-// CHECK: func.func private @_mlir_apfloat_subtract(i32, i64, i64) -> i64
+// CHECK: func.func private @__mlir_apfloat_subtract(i32, i64, i64) -> i64
 // CHECK: %[[sem:.*]] = arith.constant 18 : i32
-// CHECK: call @_mlir_apfloat_subtract(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+// CHECK: call @__mlir_apfloat_subtract(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
 func.func @subf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
   %0 = arith.subf %arg0, %arg1 : f4E2M1FN
   return
 }
 
 // -----
 
-// CHECK: func.func private @_mlir_apfloat_multiply(i32, i64, i64) -> i64
+// CHECK: func.func private @__mlir_apfloat_multiply(i32, i64, i64) -> i64
 // CHECK: %[[sem:.*]] = arith.constant 18 : i32
-// CHECK: call @_mlir_apfloat_multiply(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+// CHECK: call @__mlir_apfloat_multiply(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
 func.func @subf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
   %0 = arith.mulf %arg0, %arg1 : f4E2M1FN
   return
 }
 
 // -----
 
-// CHECK: func.func private @_mlir_apfloat_divide(i32, i64, i64) -> i64
+// CHECK: func.func private @__mlir_apfloat_divide(i32, i64, i64) -> i64
 // CHECK: %[[sem:.*]] = arith.constant 18 : i32
-// CHECK: call @_mlir_apfloat_divide(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+// CHECK: call @__mlir_apfloat_divide(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
 func.func @subf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
   %0 = arith.divf %arg0, %arg1 : f4E2M1FN
   return
 }
 
 // -----
 
-// CHECK: func.func private @_mlir_apfloat_remainder(i32, i64, i64) -> i64
+// CHECK: func.func private @__mlir_apfloat_remainder(i32, i64, i64) -> i64
 // CHECK: %[[sem:.*]] = arith.constant 18 : i32
-// CHECK: call @_mlir_apfloat_remainder(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
+// CHECK: call @__mlir_apfloat_remainder(%[[sem]], %{{.*}}, %{{.*}}) : (i32, i64, i64) -> i64
 func.func @remf(%arg0: f4E2M1FN, %arg1: f4E2M1FN) {
   %0 = arith.remf %arg0, %arg1 : f4E2M1FN
   return

Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,6 @@ def ArithToAPFloatConversionPass`
`199`	`199`	`floating-point arithmetic operations.`
`200`	`200`	`}];`
`201`	`201`	`let dependentDialects = ["func::FuncDialect"];`
`202`		`- let options = [];`
`203`	`202`	`}`
`204`	`203`
`205`	`204`	`//===----------------------------------------------------------------------===//`