[MLIR][Python] Add structured.fuseop to generator. #120601

nujaa · 2024-12-19T16:27:07Z

Implements a python interface for structured fuseOp allowing more freedom on inputs.

llvmbot · 2024-12-19T16:27:56Z

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Author: Hugo Trachino (nujaa)

Changes

Implements a python interface for structured fuseOp allowing more freedom on inputs.

Full diff: https://github.com/llvm/llvm-project/pull/120601.diff

5 Files Affected:

(modified) mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h (+1-1)
(modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp (+3-3)
(modified) mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp (+3-3)
(modified) mlir/python/mlir/dialects/transform/structured.py (+52)
(modified) mlir/test/python/dialects/transform_structured_ext.py (+21)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index 90021ffa7c380bc..efbe5c56a219b37 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -162,7 +162,7 @@ struct SparsifierOptions : public PassPipelineOptions<SparsifierOptions> {
   }
 
   /// Projects out the options for `createConvertVectorToLLVMPass`.
-  ConvertVectorToLLVMPassOptions lowerVectorToLLVMOptions() const {
+  ConvertVectorToLLVMPassOptions convertVectorToLLVMOptions() const {
     ConvertVectorToLLVMPassOptions opts{};
     opts.reassociateFPReductions = reassociateFPReductions;
     opts.force32BitVectorIndices = force32BitVectorIndices;
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
index 55143d5939ba257..842d239cf6a512c 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
@@ -35,8 +35,8 @@ using namespace mlir;
 using namespace mlir::vector;
 
 namespace {
-struct LowerVectorToLLVMPass
-    : public impl::ConvertVectorToLLVMPassBase<LowerVectorToLLVMPass> {
+struct ConvertVectorToLLVMPass
+    : public impl::ConvertVectorToLLVMPassBase<ConvertVectorToLLVMPass> {
 
   using Base::Base;
 
@@ -58,7 +58,7 @@ struct LowerVectorToLLVMPass
 };
 } // namespace
 
-void LowerVectorToLLVMPass::runOnOperation() {
+void ConvertVectorToLLVMPass::runOnOperation() {
   // Perform progressive lowering of operations on slices and
   // all contraction operations. Also applies folding and DCE.
   {
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index c5eb965884396ae..5e49252c0e57d98 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -76,16 +76,16 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
   pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass());
   pm.addPass(memref::createExpandStridedMetadataPass());
   pm.addPass(createLowerAffinePass());
-  pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
+  pm.addPass(createConvertVectorToLLVMPass(options.convertVectorToLLVMOptions()));
   pm.addPass(createFinalizeMemRefToLLVMConversionPass());
   pm.addNestedPass<func::FuncOp>(createConvertComplexToStandardPass());
   pm.addNestedPass<func::FuncOp>(arith::createArithExpandOpsPass());
   pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass());
   pm.addPass(createConvertMathToLibmPass());
   pm.addPass(createConvertComplexToLibmPass());
-  pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
+  pm.addPass(createConvertVectorToLLVMPass(options.convertVectorToLLVMOptions()));
   pm.addPass(createConvertComplexToLLVMPass());
-  pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
+  pm.addPass(createConvertVectorToLLVMPass(options.convertVectorToLLVMOptions()));
   pm.addPass(createConvertFuncToLLVMPass());
 
   // Finalize GPU code generation.
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 41051c0d5b2ffb6..b97a1aa8a82bfcf 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,
         )
 
+@_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 3ea73e8beea3688..551c2fa1e48acd3 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)
 
 
+@run
+@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)
+
+
+@run
+@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):

github-actions · 2024-12-19T16:32:16Z

✅ With the latest revision this PR passed the Python code formatter.

makslevental

this is probably good but I'm not a transform user so we'll have to wait for @ftynse to comment on the API

ftynse

LGTM with some nits.

ftynse · 2024-12-20T08:38:03Z

mlir/python/mlir/dialects/transform/structured.py

Suggested change

num_loops = sum(v if v == 0 else 1 for v in sizes)

num_loops = sum(0 if v == 0 else 1 for v in sizes)

This is slightly easier to read.

ftynse · 2024-12-20T08:39:14Z

mlir/python/mlir/dialects/transform/structured.py

Suggested change

loop_types_or_target: Union[Type, List[Type], Operation, Value],

loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView, Value],

ftynse · 2024-12-20T08:44:24Z

mlir/python/mlir/dialects/transform/structured.py

I think you need to explicitly list both overloads. And use exactly the same names for arguments, otherwise the linter complains.

Hi, sorry for the long silence, What do you mean by you need to explicitly list both overloads ? Do you mean adding a @overload under ... ? I followed what was done in other examples. The last overload one never has @overload.

Renamed arguments to tile_sizes and tile_interchange

AFAIU, the code needs to list overloads separately from the implementation, the code currently lists only one overload. https://mypy.readthedocs.io/en/stable/more_types.html#function-overloading

Added an overload setting output types.

ftynse · 2024-12-20T08:45:29Z

mlir/test/python/dialects/transform_structured_ext.py

Nit: can we add a test where sizes are not constant?

Added support for arrayAttr.

nujaa requested review from PeimingLiu, aartbik, banach-space, dcaballe, ftynse, makslevental, matthias-springer, nicolasvasilache, stellaraccident and yinying-lisa-li as code owners December 19, 2024 16:27

llvmbot added mlir:sparse Sparse compiler in MLIR mlir:python MLIR Python bindings mlir labels Dec 19, 2024

nujaa force-pushed the hugo.fuseopPython branch from 8a4bf10 to dd7406b Compare December 19, 2024 16:28

makslevental reviewed Dec 19, 2024

View reviewed changes

ftynse approved these changes Dec 20, 2024

View reviewed changes

nujaa added 4 commits January 2, 2025 19:57

[MLIR][Python] Add structured.fuseop to generator.

5849a84

fixup! [MLIR][Python] Add structured.fuseop to generator.

84cf157

fixup! fixup! [MLIR][Python] Add structured.fuseop to generator.

9e10ee1

fixup! fixup! fixup! [MLIR][Python] Add structured.fuseop to generator.

9379d82

nujaa force-pushed the hugo.fuseopPython branch from e75878a to 9379d82 Compare January 2, 2025 11:57

nujaa added 5 commits January 2, 2025 22:14

Add support for apply cleanup

be5b0cc

fixup! Add support for apply cleanup

c4f1280

fixup! fixup! Add support for apply cleanup

05e102d

Add overload initializing output type

0fd7e7d

fixup! Add overload initializing output type

3244eac

ftynse approved these changes Jan 3, 2025

View reviewed changes

nujaa merged commit 579ced4 into llvm:main Jan 3, 2025
8 checks passed

	num_loops = sum(v if v == 0 else 1 for v in sizes)
	num_loops = sum(0 if v == 0 else 1 for v in sizes)

	loop_types_or_target: Union[Type, List[Type], Operation, Value],
	loop_types_or_target: Union[Type, Sequence[Type], Operation, OpView, Value],

[MLIR][Python] Add structured.fuseop to generator. #120601

[MLIR][Python] Add structured.fuseop to generator. #120601

Uh oh!

Conversation

nujaa commented Dec 19, 2024

Uh oh!

llvmbot commented Dec 19, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Dec 19, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

makslevental left a comment

Choose a reason for hiding this comment

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

llvmbot commented Dec 19, 2024 •

edited

Loading

github-actions bot commented Dec 19, 2024 •

edited

Loading