[Mlir-commits] [mlir] fee0804 - [mlir][math] Expose math patterns as transforms (#230042)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 01:07:25 PDT 2026
Author: Adam Siemieniuk
Date: 2026-10-15T10:07:19+02:00
New Revision: fee08048551c6b9232a27411bc38b15eac4250b7
URL: https://github.com/llvm/llvm-project/commit/fee08048551c6b9232a27411bc38b15eac4250b7
DIFF: https://github.com/llvm/llvm-project/commit/fee08048551c6b9232a27411bc38b15eac4250b7.diff
LOG: [mlir][math] Expose math patterns as transforms (#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
Added:
mlir/include/mlir/Dialect/Math/TransformOps/CMakeLists.txt
mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.h
mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
mlir/python/mlir/dialects/MathTransformOps.td
mlir/python/mlir/dialects/transform/math.py
mlir/test/python/dialects/transform_math_ext.py
Modified:
mlir/include/mlir/Dialect/Math/CMakeLists.txt
mlir/lib/Dialect/Math/CMakeLists.txt
mlir/lib/RegisterAllExtensions.cpp
mlir/python/CMakeLists.txt
mlir/test/Dialect/Math/polynomial-approximation.mlir
mlir/test/lib/Dialect/Math/CMakeLists.txt
mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
mlir/test/mlir-runner/math-polynomial-approx.mlir
mlir/tools/mlir-opt/mlir-opt.cpp
Removed:
mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp
################################################################################
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 64554f0bac14a4..d7a57a15f82420 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