[Mlir-commits] [mlir] [mlir][math] Expose math patterns as transforms (PR #230042)
Adam Siemieniuk
llvmlistbot at llvm.org
Thu Oct 8 02:37:06 PDT 2026
https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/230042
Adds math dialect transform ops to start exposing individual patterns.
Initially, f32 expansion and polynomial approximation patterns are added.
The existing approximation test is converted from a test pass to the new transform ops to unify usage.
Assisted-by: Claude
>From a98809b41471c22b26aa9a44a65939315f63cfd4 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Wed, 7 Oct 2026 11:12:45 -0700
Subject: [PATCH] [mlir][math] Expose math patterns as transforms
Adds math dialect transform ops to start exposing individual patterns.
Initially, f32 expansion and polynomial approximation patterns are
added. The existing approximation test is converted from a test pass
to use the new transform ops to unify usage.
---
mlir/include/mlir/Dialect/Math/CMakeLists.txt | 1 +
.../Dialect/Math/TransformOps/CMakeLists.txt | 6 ++
.../Math/TransformOps/MathTransformOps.h | 30 ++++++++
.../Math/TransformOps/MathTransformOps.td | 39 ++++++++++
mlir/lib/Dialect/Math/CMakeLists.txt | 1 +
.../Dialect/Math/TransformOps/CMakeLists.txt | 18 +++++
.../Math/TransformOps/MathTransformOps.cpp | 67 +++++++++++++++++
mlir/lib/RegisterAllExtensions.cpp | 2 +
mlir/python/CMakeLists.txt | 9 +++
mlir/python/mlir/dialects/MathTransformOps.td | 14 ++++
mlir/python/mlir/dialects/transform/math.py | 5 ++
.../Math/polynomial-approximation.mlir | 24 ++++++-
mlir/test/lib/Dialect/Math/CMakeLists.txt | 2 -
.../Math/TestPolynomialApproximation.cpp | 71 -------------------
.../X86/math-polynomial-approx-avx2.mlir | 13 +++-
.../mlir-runner/math-polynomial-approx.mlir | 12 +++-
.../python/dialects/transform_math_ext.py | 48 +++++++++++++
mlir/tools/mlir-opt/mlir-opt.cpp | 2 -
18 files changed, 285 insertions(+), 79 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Math/TransformOps/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.h
create mode 100644 mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
create mode 100644 mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
create mode 100644 mlir/python/mlir/dialects/MathTransformOps.td
create mode 100644 mlir/python/mlir/dialects/transform/math.py
delete mode 100644 mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp
create mode 100644 mlir/test/python/dialects/transform_math_ext.py
diff --git a/mlir/include/mlir/Dialect/Math/CMakeLists.txt b/mlir/include/mlir/Dialect/Math/CMakeLists.txt
index 9f57627c321fb0..660deb21479d29 100644
--- a/mlir/include/mlir/Dialect/Math/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Math/CMakeLists.txt
@@ -1,2 +1,3 @@
add_subdirectory(IR)
+add_subdirectory(TransformOps)
add_subdirectory(Transforms)
diff --git a/mlir/include/mlir/Dialect/Math/TransformOps/CMakeLists.txt b/mlir/include/mlir/Dialect/Math/TransformOps/CMakeLists.txt
new file mode 100644
index 00000000000000..851de500d6b16c
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Math/TransformOps/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS MathTransformOps.td)
+mlir_tablegen(MathTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(MathTransformOps.cpp.inc -gen-op-defs)
+add_mlir_dialect_tablegen_target(MLIRMathTransformOpsIncGen)
+
+add_mlir_doc(MathTransformOps MathTransformOps Dialects/ -gen-op-doc -dialect=transform)
diff --git a/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.h b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.h
new file mode 100644
index 00000000000000..0a1a113f75aa93
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.h
@@ -0,0 +1,30 @@
+//===- MathTransformOps.h - Math transform ops ------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MATH_TRANSFORMOPS_MATHTRANSFORMOPS_H
+#define MLIR_DIALECT_MATH_TRANSFORMOPS_MATHTRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+#include "mlir/IR/OpImplementation.h"
+
+//===----------------------------------------------------------------------===//
+// Math Transform Operations
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Math/TransformOps/MathTransformOps.h.inc"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace math {
+void registerTransformDialectExtension(DialectRegistry ®istry);
+} // namespace math
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MATH_TRANSFORMOPS_MATHTRANSFORMOPS_H
diff --git a/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
new file mode 100644
index 00000000000000..32ed7416eb6046
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
@@ -0,0 +1,39 @@
+//===- MathTransformOps.td - Math transform ops ------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MATH_TRANSFORM_OPS
+#define MATH_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def ApplyF32ExpansionPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.math.f32_expansion",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Converts math functions to f32: operands are extended and the result is
+ truncated back.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
+def ApplyPolynomialApproximationPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.math.polynomial_approximation",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Enables polynomial approximations for math functions.
+ Optionally, allows usage of x86 AVX2 intrinsics.
+ }];
+
+ let arguments = (ins UnitAttr:$enable_avx2);
+ let assemblyFormat = "(`enable_avx2` $enable_avx2^)? attr-dict";
+}
+
+#endif // MATH_TRANSFORM_OPS
diff --git a/mlir/lib/Dialect/Math/CMakeLists.txt b/mlir/lib/Dialect/Math/CMakeLists.txt
index 9f57627c321fb0..660deb21479d29 100644
--- a/mlir/lib/Dialect/Math/CMakeLists.txt
+++ b/mlir/lib/Dialect/Math/CMakeLists.txt
@@ -1,2 +1,3 @@
add_subdirectory(IR)
+add_subdirectory(TransformOps)
add_subdirectory(Transforms)
diff --git a/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
new file mode 100644
index 00000000000000..c10e0ce55801ec
--- /dev/null
+++ b/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
@@ -0,0 +1,18 @@
+add_mlir_dialect_library(MLIRMathTransformOps
+ MathTransformOps.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Math/TransformOps
+
+ DEPENDS
+ MLIRMathTransformOpsIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRIR
+ MLIRMathDialect
+ MLIRMathTransforms
+ MLIRTransformDialect
+ MLIRVectorDialect
+ MLIRX86Dialect
+)
diff --git a/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp b/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
new file mode 100644
index 00000000000000..37fc18a6b5da95
--- /dev/null
+++ b/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
@@ -0,0 +1,67 @@
+//===- MathTransformOps.cpp - Implementation of Math transform ops --------===//
+//
+// 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/Math/TransformOps/MathTransformOps.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/X86/X86Dialect.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Apply...PatternsOp
+//===----------------------------------------------------------------------===//
+
+void transform::ApplyF32ExpansionPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ populateMathF32ExpansionPatterns(patterns, [](StringRef) { return true; });
+}
+
+void transform::ApplyPolynomialApproximationPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ bool enableAvx2 = getEnableAvx2();
+ populateMathPolynomialApproximationPatterns(patterns, [&](StringRef name) {
+ // `rsqrt` requires AVX2 so always gate it.
+ return enableAvx2 || name != math::RsqrtOp::getOperationName();
+ });
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class MathTransformDialectExtension
+ : public transform::TransformDialectExtension<
+ MathTransformDialectExtension> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MathTransformDialectExtension)
+
+ MathTransformDialectExtension() {
+ declareGeneratedDialect<arith::ArithDialect>();
+ declareGeneratedDialect<math::MathDialect>();
+ declareGeneratedDialect<vector::VectorDialect>();
+ declareGeneratedDialect<x86::X86Dialect>();
+ registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Math/TransformOps/MathTransformOps.cpp.inc"
+ >();
+ }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Math/TransformOps/MathTransformOps.cpp.inc"
+
+void mlir::math::registerTransformDialectExtension(DialectRegistry ®istry) {
+ registry.addExtensions<MathTransformDialectExtension>();
+}
diff --git a/mlir/lib/RegisterAllExtensions.cpp b/mlir/lib/RegisterAllExtensions.cpp
index e32e1142150925..67df66c0136096 100644
--- a/mlir/lib/RegisterAllExtensions.cpp
+++ b/mlir/lib/RegisterAllExtensions.cpp
@@ -43,6 +43,7 @@
#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h"
#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
+#include "mlir/Dialect/Math/TransformOps/MathTransformOps.h"
#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h"
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
@@ -104,6 +105,7 @@ void mlir::registerAllExtensions(DialectRegistry ®istry) {
func::registerTransformDialectExtension(registry);
gpu::registerTransformDialectExtension(registry);
linalg::registerTransformDialectExtension(registry);
+ math::registerTransformDialectExtension(registry);
memref::registerTransformDialectExtension(registry);
nvgpu::registerTransformDialectExtension(registry);
scf::registerTransformDialectExtension(registry);
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index aa0adea66d2fed..ca1102429ac54b 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -324,6 +324,15 @@ declare_mlir_dialect_extension_python_bindings(
DIALECT_NAME transform
EXTENSION_NAME tensor_transform)
+declare_mlir_dialect_extension_python_bindings(
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/MathTransformOps.td
+ SOURCES
+ dialects/transform/math.py
+ DIALECT_NAME transform
+ EXTENSION_NAME math_transform)
+
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/MathTransformOps.td b/mlir/python/mlir/dialects/MathTransformOps.td
new file mode 100644
index 00000000000000..8df7fd3523ef48
--- /dev/null
+++ b/mlir/python/mlir/dialects/MathTransformOps.td
@@ -0,0 +1,14 @@
+//===-- MathTransformOps.td --------------------------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_MATHTRANSFORMOPS
+#define PYTHON_BINDINGS_MATHTRANSFORMOPS
+
+include "mlir/Dialect/Math/TransformOps/MathTransformOps.td"
+
+#endif // PYTHON_BINDINGS_MATHTRANSFORMOPS
diff --git a/mlir/python/mlir/dialects/transform/math.py b/mlir/python/mlir/dialects/transform/math.py
new file mode 100644
index 00000000000000..01377d2b361508
--- /dev/null
+++ b/mlir/python/mlir/dialects/transform/math.py
@@ -0,0 +1,5 @@
+# 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
+
+from .._math_transform_ops_gen import *
diff --git a/mlir/test/Dialect/Math/polynomial-approximation.mlir b/mlir/test/Dialect/Math/polynomial-approximation.mlir
index c14ed47853411e..41c7a2819e6e85 100644
--- a/mlir/test/Dialect/Math/polynomial-approximation.mlir
+++ b/mlir/test/Dialect/Math/polynomial-approximation.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -test-math-polynomial-approximation | FileCheck %s
-// RUN: mlir-opt %s -test-math-polynomial-approximation=enable-avx2 \
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter=entry-point=avx2 \
// RUN: | FileCheck --check-prefix=AVX2 %s
// Check that all math functions lowered to approximations built from
@@ -1052,3 +1052,23 @@ func.func @rsqrt_f16(%arg0 : vector<2x8xf16>) -> vector<2x8xf16> {
%0 = "math.rsqrt"(%arg0) : (vector<2x8xf16>) -> vector<2x8xf16>
return %0 : vector<2x8xf16>
}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.math.f32_expansion
+ transform.apply_patterns.math.polynomial_approximation
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @avx2(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.math.f32_expansion
+ transform.apply_patterns.math.polynomial_approximation enable_avx2
+ } : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/lib/Dialect/Math/CMakeLists.txt b/mlir/test/lib/Dialect/Math/CMakeLists.txt
index 5d4724b9be985a..98def18df99eec 100644
--- a/mlir/test/lib/Dialect/Math/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Math/CMakeLists.txt
@@ -1,7 +1,6 @@
# Exclude tests from libMLIR.so
add_mlir_library(MLIRMathTestPasses
TestAlgebraicSimplification.cpp
- TestPolynomialApproximation.cpp
EXCLUDE_FROM_LIBMLIR
)
@@ -10,5 +9,4 @@ mlir_target_link_libraries(MLIRMathTestPasses PUBLIC
MLIRPass
MLIRTransformUtils
MLIRVectorDialect
- MLIRX86Dialect
)
diff --git a/mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp b/mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp
deleted file mode 100644
index 486ccc0c307655..00000000000000
--- a/mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp
+++ /dev/null
@@ -1,71 +0,0 @@
-//===- TestPolynomialApproximation.cpp - Test math ops approximations -----===//
-//
-// 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 contains test passes for expanding math operations into
-// polynomial approximations.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Math/IR/Math.h"
-#include "mlir/Dialect/Math/Transforms/Passes.h"
-#include "mlir/Dialect/Vector/IR/VectorDialect.h"
-#include "mlir/Dialect/X86/X86Dialect.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-using namespace mlir;
-
-namespace {
-struct TestMathPolynomialApproximationPass
- : public PassWrapper<TestMathPolynomialApproximationPass, OperationPass<>> {
- MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
- TestMathPolynomialApproximationPass)
-
- TestMathPolynomialApproximationPass() = default;
- TestMathPolynomialApproximationPass(
- const TestMathPolynomialApproximationPass &pass)
- : PassWrapper(pass) {}
-
- void runOnOperation() override;
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, math::MathDialect,
- vector::VectorDialect>();
- if (enableAvx2)
- registry.insert<x86::X86Dialect>();
- }
- StringRef getArgument() const final {
- return "test-math-polynomial-approximation";
- }
- StringRef getDescription() const final {
- return "Test math polynomial approximations";
- }
-
- Option<bool> enableAvx2{
- *this, "enable-avx2",
- llvm::cl::desc("Enable approximations that emit AVX2 intrinsics via the "
- "X86 dialect"),
- llvm::cl::init(false)};
-};
-} // namespace
-
-void TestMathPolynomialApproximationPass::runOnOperation() {
- RewritePatternSet patterns(&getContext());
- MathPolynomialApproximationOptions approxOptions;
- approxOptions.enableAvx2 = enableAvx2;
- populateMathPolynomialApproximationPatterns(patterns, approxOptions);
- (void)applyPatternsGreedily(getOperation(), std::move(patterns));
-}
-
-namespace mlir {
-namespace test {
-void registerTestMathPolynomialApproximationPass() {
- PassRegistration<TestMathPolynomialApproximationPass>();
-}
-} // namespace test
-} // namespace mlir
diff --git a/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir b/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
index 16b363de4d6304..6a1bc237c00a89 100644
--- a/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
+++ b/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
@@ -1,4 +1,5 @@
-// RUN: mlir-opt %s -test-math-polynomial-approximation="enable-avx2" \
+// RUN: mlir-opt %s -transform-interpreter \
+// RUN: -test-transform-dialect-erase-schedule \
// RUN: -convert-vector-to-scf \
// RUN: -convert-scf-to-cf \
// RUN: -convert-arith-to-llvm \
@@ -41,3 +42,13 @@ func.func @main() {
call @rsqrt(): () -> ()
return
}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.math.polynomial_approximation enable_avx2
+ } : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/mlir-runner/math-polynomial-approx.mlir b/mlir/test/mlir-runner/math-polynomial-approx.mlir
index 5bde6b41c3f86a..3bf412aec8df45 100644
--- a/mlir/test/mlir-runner/math-polynomial-approx.mlir
+++ b/mlir/test/mlir-runner/math-polynomial-approx.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-math-polynomial-approximation),convert-vector-to-scf,convert-scf-to-cf,convert-vector-to-llvm,convert-to-llvm,reconcile-unrealized-casts)" \
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(transform-interpreter,test-transform-dialect-erase-schedule,convert-vector-to-scf,convert-scf-to-cf,convert-vector-to-llvm,convert-to-llvm,reconcile-unrealized-casts)" \
// RUN: | mlir-runner \
// RUN: -e main -entry-point-result=void -O0 \
// RUN: -shared-libs=%mlir_c_runner_utils \
@@ -858,3 +858,13 @@ func.func @main() {
call @ceilf() : () -> ()
return
}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.math.polynomial_approximation
+ } : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/python/dialects/transform_math_ext.py b/mlir/test/python/dialects/transform_math_ext.py
new file mode 100644
index 00000000000000..418c837332ead2
--- /dev/null
+++ b/mlir/test/python/dialects/transform_math_ext.py
@@ -0,0 +1,48 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import transform
+from mlir.dialects.transform import math
+
+
+def run_apply_patterns(f):
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.AnyOpType.get(),
+ )
+ with InsertionPoint(sequence.body):
+ apply = transform.ApplyPatternsOp(sequence.bodyTarget)
+ with InsertionPoint(apply.patterns):
+ f()
+ transform.YieldOp()
+ print("\nTEST:", f.__name__)
+ print(module)
+ return f
+
+
+ at run_apply_patterns
+def testF32Expansion():
+ math.ApplyF32ExpansionPatternsOp()
+ # CHECK-LABEL: TEST: testF32Expansion
+ # CHECK: apply_patterns
+ # CHECK: transform.apply_patterns.math.f32_expansion
+
+
+ at run_apply_patterns
+def testPolynomialApproximation():
+ math.ApplyPolynomialApproximationPatternsOp()
+ # CHECK-LABEL: TEST: testPolynomialApproximation
+ # CHECK: apply_patterns
+ # CHECK: transform.apply_patterns.math.polynomial_approximation{{$}}
+
+
+ at run_apply_patterns
+def testPolynomialApproximationAvx2():
+ math.ApplyPolynomialApproximationPatternsOp(enable_avx2=True)
+ # CHECK-LABEL: TEST: testPolynomialApproximationAvx2
+ # CHECK: apply_patterns
+ # CHECK: transform.apply_patterns.math.polynomial_approximation enable_avx2
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 730e48914281a0..f4f8c5943580ec 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -125,7 +125,6 @@ void registerTestLowerToLLVM();
void registerTestMakeIsolatedFromAbovePass();
void registerTestMatchReductionPass();
void registerTestMathAlgebraicSimplificationPass();
-void registerTestMathPolynomialApproximationPass();
void registerTestMathToVCIXPass();
void registerTestIrdlTestDialectConversionPass();
void registerTestMemRefDependenceCheck();
@@ -276,7 +275,6 @@ static void registerTestPasses() {
mlir::test::registerTestMakeIsolatedFromAbovePass();
mlir::test::registerTestMatchReductionPass();
mlir::test::registerTestMathAlgebraicSimplificationPass();
- mlir::test::registerTestMathPolynomialApproximationPass();
mlir::test::registerTestMathToVCIXPass();
mlir::test::registerTestMemRefDependenceCheck();
mlir::test::registerTestMemRefStrideCalculation();
More information about the Mlir-commits
mailing list