aac4eb5c3c
In https://github.com/llvm/llvm-project/pull/94714, we add a python function `apply_patterns_and_fold_greedily` which accepts an `MlirModule` as the argument type. However, sometimes we want to apply patterns with an `MlirOperation` argument, and there is currently no python API to convert an `MlirOperation` to `MlirModule`. So here we overload this function `apply_patterns_and_fold_greedily` to do this (also a corresponding new C API `mlirApplyPatternsAndFoldGreedilyWithOp`)
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
# RUN: %PYTHON %s 2>&1 | FileCheck %s
|
|
|
|
from mlir.dialects import arith, func, pdl
|
|
from mlir.dialects.builtin import module
|
|
from mlir.ir import *
|
|
from mlir.rewrite import *
|
|
|
|
|
|
def construct_and_print_in_module(f):
|
|
print("\nTEST:", f.__name__)
|
|
with Context(), Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
module = f(module)
|
|
if module is not None:
|
|
print(module)
|
|
return f
|
|
|
|
def get_pdl_patterns():
|
|
# Create a rewrite from add to mul. This will match
|
|
# - operation name is arith.addi
|
|
# - operands are index types.
|
|
# - there are two operands.
|
|
with Location.unknown():
|
|
m = Module.create()
|
|
with InsertionPoint(m.body):
|
|
# Change all arith.addi with index types to arith.muli.
|
|
@pdl.pattern(benefit=1, sym_name="addi_to_mul")
|
|
def pat():
|
|
# Match arith.addi with index types.
|
|
index_type = pdl.TypeOp(IndexType.get())
|
|
operand0 = pdl.OperandOp(index_type)
|
|
operand1 = pdl.OperandOp(index_type)
|
|
op0 = pdl.OperationOp(
|
|
name="arith.addi", args=[operand0, operand1], types=[index_type]
|
|
)
|
|
|
|
# Replace the matched op with arith.muli.
|
|
@pdl.rewrite()
|
|
def rew():
|
|
newOp = pdl.OperationOp(
|
|
name="arith.muli", args=[operand0, operand1], types=[index_type]
|
|
)
|
|
pdl.ReplaceOp(op0, with_op=newOp)
|
|
|
|
# Create a PDL module from module and freeze it. At this point the ownership
|
|
# of the module is transferred to the PDL module. This ownership transfer is
|
|
# not yet captured Python side/has sharp edges. So best to construct the
|
|
# module and PDL module in same scope.
|
|
# FIXME: This should be made more robust.
|
|
return PDLModule(m).freeze()
|
|
|
|
|
|
# CHECK-LABEL: TEST: test_add_to_mul
|
|
# CHECK: arith.muli
|
|
@construct_and_print_in_module
|
|
def test_add_to_mul(module_):
|
|
index_type = IndexType.get()
|
|
|
|
# Create a test case.
|
|
@module(sym_name="ir")
|
|
def ir():
|
|
@func.func(index_type, index_type)
|
|
def add_func(a, b):
|
|
return arith.addi(a, b)
|
|
|
|
frozen = get_pdl_patterns()
|
|
# Could apply frozen pattern set multiple times.
|
|
apply_patterns_and_fold_greedily(module_, frozen)
|
|
return module_
|
|
|
|
|
|
# CHECK-LABEL: TEST: test_add_to_mul_with_op
|
|
# CHECK: arith.muli
|
|
@construct_and_print_in_module
|
|
def test_add_to_mul_with_op(module_):
|
|
index_type = IndexType.get()
|
|
|
|
# Create a test case.
|
|
@module(sym_name="ir")
|
|
def ir():
|
|
@func.func(index_type, index_type)
|
|
def add_func(a, b):
|
|
return arith.addi(a, b)
|
|
|
|
frozen = get_pdl_patterns()
|
|
apply_patterns_and_fold_greedily(module_.operation, frozen)
|
|
return module_
|