[Mlir-commits] [mlir] [mlir] Place MathToFuncs helpers in the nearest symbol table (PR #232582)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 20:48:27 PDT 2026
https://github.com/frokaikan created https://github.com/llvm/llvm-project/pull/232582
The MathToFuncs pass previously inserted every outlined ipowi, fpowi, and ctlz helper into the top-level builtin.module. Calls rewritten inside a gpu.module could not resolve those helpers because the declarations were outside the nearest symbol table.
Create each helper in the symbol table nearest to its source operation. Include that symbol table in the implementation cache key, and use the same key when looking up helpers during conversion. This preserves reuse for repeated calls in one module without sharing a definition across distinct gpu.modules or the host module.
Add a regression test with two gpu.modules and a host module. Cover the valid scalar ipowi, fpowi, and ctlz type combinations, two calls of each combination per module, and vector lowering. Check that every module receives exactly one private definition per helper and that all calls resolve within their own symbol table.
Tested with an incremental mlir-opt build and the four MathToFuncs lit tests (ASAN_OPTIONS=detect_leaks=0 for the sandboxed test run).
>From ad5fce7256f374148b29e5c7dd7a5b923680795d Mon Sep 17 00:00:00 2001
From: mingruiz <mingruiz at nvidia.com>
Date: Thu, 15 Oct 2026 20:36:36 -0700
Subject: [PATCH] [mlir] Place MathToFuncs helpers in the nearest symbol table
The MathToFuncs pass previously inserted every outlined ipowi, fpowi, and
ctlz helper into the top-level builtin.module. Calls rewritten inside a
gpu.module could not resolve those helpers because the declarations were
outside the nearest symbol table.
Create each helper in the symbol table nearest to its source operation.
Include that symbol table in the implementation cache key, and use the
same key when looking up helpers during conversion. This preserves reuse
for repeated calls in one module without sharing a definition across
distinct gpu.modules or the host module.
Add a regression test with two gpu.modules and a host module. Cover the
valid scalar ipowi, fpowi, and ctlz type combinations, two calls of each
combination per module, and vector lowering. Check that every module
receives exactly one private definition per helper and that all calls
resolve within their own symbol table.
Tested with an incremental mlir-opt build and the four MathToFuncs lit
tests (ASAN_OPTIONS=detect_leaks=0 for the sandboxed test run).
---
.../Conversion/MathToFuncs/MathToFuncs.cpp | 58 ++++----
.../Conversion/MathToFuncs/symbol-tables.mlir | 134 ++++++++++++++++++
2 files changed, 168 insertions(+), 24 deletions(-)
create mode 100644 mlir/test/Conversion/MathToFuncs/symbol-tables.mlir
diff --git a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
index f653c6b38e002a..3bb6c6901d2649 100644
--- a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
+++ b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
@@ -24,6 +25,8 @@
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/DebugLog.h"
+#include <tuple>
+
namespace mlir {
#define GEN_PASS_DEF_CONVERTMATHTOFUNCS
#include "mlir/Conversion/Passes.h.inc"
@@ -151,7 +154,7 @@ static FunctionType getElementalFuncTypeForOp(Operation *op) {
}
/// Create linkonce_odr function to implement the power function with
-/// the given \p elementType type inside \p module. The \p elementType
+/// the given \p elementType type inside \p symbolTable. The \p elementType
/// must be IntegerType, an the created function has
/// 'IntegerType (*)(IntegerType, IntegerType)' function type.
///
@@ -181,12 +184,13 @@ static FunctionType getElementalFuncTypeForOp(Operation *op) {
/// b *= b;
/// }
/// }
-static func::FuncOp createElementIPowIFunc(ModuleOp *module, Type elementType) {
+static func::FuncOp createElementIPowIFunc(Operation *symbolTable,
+ Type elementType) {
assert(isa<IntegerType>(elementType) &&
"non-integer element type for IPowIOp");
- ImplicitLocOpBuilder builder =
- ImplicitLocOpBuilder::atBlockEnd(module->getLoc(), module->getBody());
+ ImplicitLocOpBuilder builder = ImplicitLocOpBuilder::atBlockEnd(
+ symbolTable->getLoc(), &symbolTable->getRegion(0).front());
std::string funcName("__mlir_math_ipowi");
llvm::raw_string_ostream nameOS(funcName);
@@ -378,7 +382,7 @@ IPowIOpLowering::matchAndRewrite(math::IPowIOp op,
}
/// Create linkonce_odr function to implement the power function with
-/// the given \p funcType type inside \p module. The \p funcType must be
+/// the given \p funcType type inside \p symbolTable. The \p funcType must be
/// 'FloatType (*)(FloatType, IntegerType)' function type.
///
/// template <typename T>
@@ -410,12 +414,12 @@ IPowIOpLowering::matchAndRewrite(math::IPowIOp op,
/// }
/// return result;
/// }
-static func::FuncOp createElementFPowIFunc(ModuleOp *module,
+static func::FuncOp createElementFPowIFunc(Operation *symbolTable,
FunctionType funcType) {
auto baseType = cast<FloatType>(funcType.getInput(0));
auto powType = cast<IntegerType>(funcType.getInput(1));
- ImplicitLocOpBuilder builder =
- ImplicitLocOpBuilder::atBlockEnd(module->getLoc(), module->getBody());
+ ImplicitLocOpBuilder builder = ImplicitLocOpBuilder::atBlockEnd(
+ symbolTable->getLoc(), &symbolTable->getRegion(0).front());
std::string funcName("__mlir_math_fpowi");
llvm::raw_string_ostream nameOS(funcName);
@@ -605,8 +609,8 @@ FPowIOpLowering::matchAndRewrite(math::FPowIOp op,
}
/// Create function to implement the ctlz function the given \p elementType type
-/// inside \p module. The \p elementType must be IntegerType, an the created
-/// function has 'IntegerType (*)(IntegerType)' function type.
+/// inside \p symbolTable. The \p elementType must be IntegerType, an the
+/// created function has 'IntegerType (*)(IntegerType)' function type.
///
/// template <typename T>
/// T __mlir_math_ctlz_*(T x) {
@@ -651,7 +655,7 @@ FPowIOpLowering::matchAndRewrite(math::FPowIOp op,
/// }
/// return %out: i32
/// }
-static func::FuncOp createCtlzFunc(ModuleOp *module, Type elementType) {
+static func::FuncOp createCtlzFunc(Operation *symbolTable, Type elementType) {
if (!isa<IntegerType>(elementType)) {
LDBG() << "non-integer element type for CtlzFunc; type was: "
<< elementType;
@@ -659,9 +663,9 @@ static func::FuncOp createCtlzFunc(ModuleOp *module, Type elementType) {
}
int64_t bitWidth = elementType.getIntOrFloatBitWidth();
- Location loc = module->getLoc();
+ Location loc = symbolTable->getLoc();
ImplicitLocOpBuilder builder =
- ImplicitLocOpBuilder::atBlockEnd(loc, module->getBody());
+ ImplicitLocOpBuilder::atBlockEnd(loc, &symbolTable->getRegion(0).front());
std::string funcName("__mlir_math_ctlz");
llvm::raw_string_ostream nameOS(funcName);
@@ -791,10 +795,10 @@ struct ConvertMathToFuncsPass
// and store them in funcImpls map.
void generateOpImplementations();
- // A map between pairs of (operation, type) deduced from operations that this
- // pass will convert, and the corresponding outlined software implementations
- // of these operations for the given type.
- DenseMap<std::pair<OperationName, Type>, func::FuncOp> funcImpls;
+ // A map between tuples of (operation, type, symbol table) and the
+ // corresponding outlined software implementations.
+ DenseMap<std::tuple<OperationName, Type, Operation *>, func::FuncOp>
+ funcImpls;
};
} // namespace
@@ -820,10 +824,11 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
// Generate the software implementation of this operation,
// if it has not been generated yet.
- auto key = std::pair(op->getName(), resultType);
+ Operation *symbolTable = SymbolTable::getNearestSymbolTable(op);
+ auto key = std::make_tuple(op->getName(), resultType, symbolTable);
auto entry = funcImpls.try_emplace(key, func::FuncOp{});
if (entry.second)
- entry.first->second = createCtlzFunc(&module, resultType);
+ entry.first->second = createCtlzFunc(symbolTable, resultType);
})
.Case([&](math::IPowIOp op) {
if (!isConvertible(op))
@@ -833,10 +838,12 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
// Generate the software implementation of this operation,
// if it has not been generated yet.
- auto key = std::pair(op->getName(), resultType);
+ Operation *symbolTable = SymbolTable::getNearestSymbolTable(op);
+ auto key = std::make_tuple(op->getName(), resultType, symbolTable);
auto entry = funcImpls.try_emplace(key, func::FuncOp{});
if (entry.second)
- entry.first->second = createElementIPowIFunc(&module, resultType);
+ entry.first->second =
+ createElementIPowIFunc(symbolTable, resultType);
})
.Case([&](math::FPowIOp op) {
if (!isFPowIConvertible(op))
@@ -848,10 +855,11 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
// if it has not been generated yet.
// FPowI implementations are mapped via the FunctionType
// created from the operation's result and operands.
- auto key = std::pair(op->getName(), funcType);
+ Operation *symbolTable = SymbolTable::getNearestSymbolTable(op);
+ auto key = std::make_tuple(op->getName(), funcType, symbolTable);
auto entry = funcImpls.try_emplace(key, func::FuncOp{});
if (entry.second)
- entry.first->second = createElementFPowIFunc(&module, funcType);
+ entry.first->second = createElementFPowIFunc(symbolTable, funcType);
});
});
}
@@ -869,7 +877,9 @@ void ConvertMathToFuncsPass::runOnOperation() {
// For the given Type Returns FuncOp stored in funcImpls map.
auto getFuncOpByType = [&](Operation *op, Type type) -> func::FuncOp {
- auto it = funcImpls.find(std::pair(op->getName(), type));
+ auto key = std::make_tuple(op->getName(), type,
+ SymbolTable::getNearestSymbolTable(op));
+ auto it = funcImpls.find(key);
if (it == funcImpls.end())
return {};
diff --git a/mlir/test/Conversion/MathToFuncs/symbol-tables.mlir b/mlir/test/Conversion/MathToFuncs/symbol-tables.mlir
new file mode 100644
index 00000000000000..c947f32a7c822e
--- /dev/null
+++ b/mlir/test/Conversion/MathToFuncs/symbol-tables.mlir
@@ -0,0 +1,134 @@
+// RUN: mlir-opt %s -verify-each -pass-pipeline="builtin.module(convert-math-to-funcs{convert-ctlz})" | FileCheck %s
+
+// CHECK-LABEL: gpu.module @first {
+// CHECK-LABEL: gpu.func @kernel(
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i64
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i32
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i64
+// CHECK-LABEL: gpu.func @kernel_vector(
+// CHECK: vector.extract
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-8: func.func private @__mlir_math_
+// CHECK-NOT: func.func private @__mlir_math_
+// CHECK-LABEL: gpu.module @second {
+// CHECK-LABEL: gpu.func @kernel(
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i64
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i32
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i64
+// CHECK-LABEL: gpu.func @kernel_vector(
+// CHECK: vector.extract
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-8: func.func private @__mlir_math_
+// CHECK-NOT: func.func private @__mlir_math_
+// CHECK-LABEL: func.func @host(
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i64
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f64_i64
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i32
+// CHECK-COUNT-2: call @__mlir_math_ctlz_i64
+// CHECK-LABEL: func.func @host_vector(
+// CHECK: vector.extract
+// CHECK-COUNT-2: call @__mlir_math_ipowi_i32
+// CHECK-COUNT-2: call @__mlir_math_fpowi_f32_i32
+// CHECK-COUNT-8: func.func private @__mlir_math_
+// CHECK-NOT: func.func private @__mlir_math_
+
+module {
+ gpu.module @first {
+ gpu.func @kernel(%base32: i32, %exp32: i32, %base64: i64,
+ %exp64: i64, %basef32: f32, %basef64: f64) kernel {
+ %ipowi32_0 = math.ipowi %base32, %exp32 : i32
+ %ipowi32_1 = math.ipowi %base32, %exp32 : i32
+ %ipowi64_0 = math.ipowi %base64, %exp64 : i64
+ %ipowi64_1 = math.ipowi %base64, %exp64 : i64
+ %fpowi_f32_i32_0 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i32_1 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i64_0 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f32_i64_1 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f64_i32_0 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i32_1 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i64_0 = math.fpowi %basef64, %exp64 : f64, i64
+ %fpowi_f64_i64_1 = math.fpowi %basef64, %exp64 : f64, i64
+ %ctlz32_0 = math.ctlz %base32 : i32
+ %ctlz32_1 = math.ctlz %base32 : i32
+ %ctlz64_0 = math.ctlz %base64 : i64
+ %ctlz64_1 = math.ctlz %base64 : i64
+ gpu.return
+ }
+ gpu.func @kernel_vector(%base32: vector<2xi32>, %exp32: vector<2xi32>,
+ %basef32: vector<2xf32>) kernel {
+ %ipowi = math.ipowi %base32, %exp32 : vector<2xi32>
+ %fpowi = math.fpowi %basef32, %exp32 : vector<2xf32>, vector<2xi32>
+ gpu.return
+ }
+ }
+ gpu.module @second {
+ gpu.func @kernel(%base32: i32, %exp32: i32, %base64: i64,
+ %exp64: i64, %basef32: f32, %basef64: f64) kernel {
+ %ipowi32_0 = math.ipowi %base32, %exp32 : i32
+ %ipowi32_1 = math.ipowi %base32, %exp32 : i32
+ %ipowi64_0 = math.ipowi %base64, %exp64 : i64
+ %ipowi64_1 = math.ipowi %base64, %exp64 : i64
+ %fpowi_f32_i32_0 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i32_1 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i64_0 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f32_i64_1 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f64_i32_0 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i32_1 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i64_0 = math.fpowi %basef64, %exp64 : f64, i64
+ %fpowi_f64_i64_1 = math.fpowi %basef64, %exp64 : f64, i64
+ %ctlz32_0 = math.ctlz %base32 : i32
+ %ctlz32_1 = math.ctlz %base32 : i32
+ %ctlz64_0 = math.ctlz %base64 : i64
+ %ctlz64_1 = math.ctlz %base64 : i64
+ gpu.return
+ }
+ gpu.func @kernel_vector(%base32: vector<2xi32>, %exp32: vector<2xi32>,
+ %basef32: vector<2xf32>) kernel {
+ %ipowi = math.ipowi %base32, %exp32 : vector<2xi32>
+ %fpowi = math.fpowi %basef32, %exp32 : vector<2xf32>, vector<2xi32>
+ gpu.return
+ }
+ }
+ func.func @host(%base32: i32, %exp32: i32, %base64: i64,
+ %exp64: i64, %basef32: f32, %basef64: f64) {
+ %ipowi32_0 = math.ipowi %base32, %exp32 : i32
+ %ipowi32_1 = math.ipowi %base32, %exp32 : i32
+ %ipowi64_0 = math.ipowi %base64, %exp64 : i64
+ %ipowi64_1 = math.ipowi %base64, %exp64 : i64
+ %fpowi_f32_i32_0 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i32_1 = math.fpowi %basef32, %exp32 : f32, i32
+ %fpowi_f32_i64_0 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f32_i64_1 = math.fpowi %basef32, %exp64 : f32, i64
+ %fpowi_f64_i32_0 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i32_1 = math.fpowi %basef64, %exp32 : f64, i32
+ %fpowi_f64_i64_0 = math.fpowi %basef64, %exp64 : f64, i64
+ %fpowi_f64_i64_1 = math.fpowi %basef64, %exp64 : f64, i64
+ %ctlz32_0 = math.ctlz %base32 : i32
+ %ctlz32_1 = math.ctlz %base32 : i32
+ %ctlz64_0 = math.ctlz %base64 : i64
+ %ctlz64_1 = math.ctlz %base64 : i64
+ func.return
+ }
+ func.func @host_vector(%base32: vector<2xi32>, %exp32: vector<2xi32>,
+ %basef32: vector<2xf32>) {
+ %ipowi = math.ipowi %base32, %exp32 : vector<2xi32>
+ %fpowi = math.fpowi %basef32, %exp32 : vector<2xf32>, vector<2xi32>
+ func.return
+ }
+}
More information about the Mlir-commits
mailing list