[Mlir-commits] [mlir] [MLIR][Python] Add structured.fuseop to generator. (PR #120601)
Hugo Trachino
llvmlistbot at llvm.org
Thu Jan 2 06:14:38 PST 2025
https://github.com/nujaa updated https://github.com/llvm/llvm-project/pull/120601
>From 5849a846f29568e9395d3aaa24fc9ebea6a97b45 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Thu, 19 Dec 2024 19:30:52 +0800
Subject: [PATCH 1/5] [MLIR][Python] Add structured.fuseop to generator.
---
.../mlir/dialects/transform/structured.py | 52 +++++++++++++++++++
.../dialects/transform_structured_ext.py | 21 ++++++++
2 files changed, 73 insertions(+)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 9121aa8e40237b..e715872a16ce48 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -139,6 +139,58 @@ def __init__(
ip=ip,
)
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class FuseOp(FuseOp):
+ """Specialization for FuseOp class."""
+
+ @overload
+ def __init__(
+ self,
+ target: Union[Operation, Value, OpView],
+ *,
+ sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+ interchange: OptionalIntList = None,
+ loc=None,
+ ip=None,
+ ):
+ ...
+
+ def __init__(
+ self,
+ loop_types_or_target: Union[Type, List[Type], Operation, Value],
+ target_or_none: Optional[Union[Operation, Value, OpView]] = None,
+ *,
+ sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+ interchange: OptionalIntList = None,
+ loc=None,
+ ip=None,
+ ):
+ sizes = sizes if sizes else []
+ num_loops = sum(v if v == 0 else 1 for v in sizes)
+
+ if isinstance(loop_types_or_target, (Operation, Value, OpView)):
+ loop_types = [transform.AnyOpType.get()] * num_loops
+ target = loop_types_or_target
+ assert (
+ target_or_none is None
+ ), "Cannot construct FuseOp with two targets."
+ else:
+ loop_types = (
+ ([loop_types_or_target] * num_loops)
+ if isinstance(loop_types_or_target, Type)
+ else loop_types_or_target
+ )
+ target = target_or_none
+ super().__init__(
+ target.type,
+ loop_types,
+ target,
+ tile_sizes=sizes,
+ tile_interchange=interchange,
+ loc=loc,
+ ip=ip,
+ )
+
@_ods_cext.register_operation(_Dialect, replace=True)
class GeneralizeOp(GeneralizeOp):
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index fb4c75b5337928..69cac29b1fbc91 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -101,6 +101,27 @@ def testFuseIntoContainingOpCompact(target):
# CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ at run
+ at create_sequence
+def testFuseOpCompact(target):
+ structured.FuseOp(target, sizes=[4, 8], interchange=[0, 1])
+ # CHECK-LABEL: TEST: testFuseOpCompact
+ # CHECK: transform.sequence
+ # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
+ # CHECK-SAME: interchange [0, 1]
+ # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+
+ at run
+ at create_sequence
+def testFuseOpNoArg(target):
+ structured.FuseOp(target)
+ # CHECK-LABEL: TEST: testFuseOpNoArg
+ # CHECK: transform.sequence
+ # CHECK: %{{.+}} = transform.structured.fuse %{{.*}} :
+ # CHECK-SAME: (!transform.any_op) -> !transform.any_op
+
+
@run
@create_sequence
def testGeneralize(target):
>From 84cf157a6dd5f48b62d1e839aeea06f3eaeb5a82 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Fri, 20 Dec 2024 00:53:27 +0800
Subject: [PATCH 2/5] fixup! [MLIR][Python] Add structured.fuseop to generator.
---
mlir/python/mlir/dialects/transform/structured.py | 4 +---
1 file changed, 1 insertion(+), 3 deletions(-)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index e715872a16ce48..e30500f512cd7b 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -171,9 +171,7 @@ def __init__(
if isinstance(loop_types_or_target, (Operation, Value, OpView)):
loop_types = [transform.AnyOpType.get()] * num_loops
target = loop_types_or_target
- assert (
- target_or_none is None
- ), "Cannot construct FuseOp with two targets."
+ assert target_or_none is None, "Cannot construct FuseOp with two targets."
else:
loop_types = (
([loop_types_or_target] * num_loops)
>From 9e10ee19b409a13ed0e2e22eec04779ce5b37559 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Fri, 20 Dec 2024 01:09:11 +0800
Subject: [PATCH 3/5] fixup! fixup! [MLIR][Python] Add structured.fuseop to
generator.
---
mlir/python/mlir/dialects/transform/structured.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index e30500f512cd7b..1d0c71639d07b4 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -139,6 +139,7 @@ def __init__(
ip=ip,
)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class FuseOp(FuseOp):
"""Specialization for FuseOp class."""
>From 9379d8280307997cb5f13628977d5df6a8f99d90 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Thu, 2 Jan 2025 19:18:06 +0800
Subject: [PATCH 4/5] fixup! fixup! fixup! [MLIR][Python] Add structured.fuseop
to generator.
---
.../mlir/dialects/transform/structured.py | 21 +++++++++++--------
.../dialects/transform_structured_ext.py | 15 ++++++++++++-
2 files changed, 26 insertions(+), 10 deletions(-)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 1d0c71639d07b4..23bbca34c618cb 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -149,8 +149,8 @@ def __init__(
self,
target: Union[Operation, Value, OpView],
*,
- sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
- interchange: OptionalIntList = None,
+ tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+ tile_interchange: OptionalIntList = None,
loc=None,
ip=None,
):
@@ -158,16 +158,19 @@ def __init__(
def __init__(
self,
- loop_types_or_target: Union[Type, List[Type], Operation, Value],
+ loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView, Value],
target_or_none: Optional[Union[Operation, Value, OpView]] = None,
*,
- sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
- interchange: OptionalIntList = None,
+ tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
+ tile_interchange: OptionalIntList = None,
loc=None,
ip=None,
):
- sizes = sizes if sizes else []
- num_loops = sum(v if v == 0 else 1 for v in sizes)
+ tile_sizes = tile_sizes if tile_sizes else []
+ tile_interchange = tile_interchange if tile_interchange else []
+ _, tile_sizes, _ = _dispatch_dynamic_index_list(tile_sizes)
+ _, tile_interchange, _ = _dispatch_dynamic_index_list(tile_interchange)
+ num_loops = sum(0 if v == 0 else 1 for v in tile_sizes)
if isinstance(loop_types_or_target, (Operation, Value, OpView)):
loop_types = [transform.AnyOpType.get()] * num_loops
@@ -184,8 +187,8 @@ def __init__(
target.type,
loop_types,
target,
- tile_sizes=sizes,
- tile_interchange=interchange,
+ tile_sizes=tile_sizes,
+ tile_interchange=tile_interchange,
loc=loc,
ip=ip,
)
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index 69cac29b1fbc91..49660d6d4c10f1 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -104,7 +104,7 @@ def testFuseIntoContainingOpCompact(target):
@run
@create_sequence
def testFuseOpCompact(target):
- structured.FuseOp(target, sizes=[4, 8], interchange=[0, 1])
+ structured.FuseOp(target, tile_sizes=[4, 8], tile_interchange=[0, 1])
# CHECK-LABEL: TEST: testFuseOpCompact
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
@@ -122,6 +122,19 @@ def testFuseOpNoArg(target):
# CHECK-SAME: (!transform.any_op) -> !transform.any_op
+ at run
+ at create_sequence
+def testFuseOpAttributes(target):
+ attr = DenseI64ArrayAttr.get([4, 8])
+ ichange = DenseI64ArrayAttr.get([0, 1])
+ structured.FuseOp(target, tile_sizes=attr, tile_interchange=ichange)
+ # CHECK-LABEL: TEST: testFuseOpAttributes
+ # CHECK: transform.sequence
+ # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
+ # CHECK-SAME: interchange [0, 1]
+ # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+
@run
@create_sequence
def testGeneralize(target):
>From be5b0cc227a67b9dc9e9d4d62c10f5fa74a24c47 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Thu, 2 Jan 2025 22:14:21 +0800
Subject: [PATCH 5/5] Add support for apply cleanup
---
mlir/python/mlir/dialects/transform/structured.py | 3 ++-
mlir/test/python/dialects/transform_structured_ext.py | 7 +++++--
2 files changed, 7 insertions(+), 3 deletions(-)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 23bbca34c618cb..b3340de59cf1ff 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -158,7 +158,8 @@ def __init__(
def __init__(
self,
- loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView, Value],
+ loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView,
+ Value],
target_or_none: Optional[Union[Operation, Value, OpView]] = None,
*,
tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index 49660d6d4c10f1..f8140725a3db50 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -104,11 +104,14 @@ def testFuseIntoContainingOpCompact(target):
@run
@create_sequence
def testFuseOpCompact(target):
- structured.FuseOp(target, tile_sizes=[4, 8], tile_interchange=[0, 1])
+ structured.FuseOp(target,
+ tile_sizes=[4, 8],
+ tile_interchange=[0, 1],
+ apply_cleanup=True)
# CHECK-LABEL: TEST: testFuseOpCompact
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
- # CHECK-SAME: interchange [0, 1]
+ # CHECK-SAME: interchange [0, 1] apply_cleanup = true
# CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
More information about the Mlir-commits
mailing list