xref: /llvm-project/mlir/test/python/dialects/transform_sparse_tensor_ext.py (revision 3d27d1152eacd6432485cd81d471bb03987a83e1)
1*3d27d115SPeiming Liu# RUN: %PYTHON %s | FileCheck %s
2*3d27d115SPeiming Liu
3*3d27d115SPeiming Liufrom mlir.ir import *
4*3d27d115SPeiming Liufrom mlir.dialects import transform
5*3d27d115SPeiming Liufrom mlir.dialects.transform import sparse_tensor
6*3d27d115SPeiming Liu
7*3d27d115SPeiming Liu
8*3d27d115SPeiming Liudef run(f):
9*3d27d115SPeiming Liu    with Context(), Location.unknown():
10*3d27d115SPeiming Liu        module = Module.create()
11*3d27d115SPeiming Liu        with InsertionPoint(module.body):
12*3d27d115SPeiming Liu            sequence = transform.SequenceOp(
13*3d27d115SPeiming Liu                transform.FailurePropagationMode.Propagate,
14*3d27d115SPeiming Liu                [],
15*3d27d115SPeiming Liu                transform.AnyOpType.get(),
16*3d27d115SPeiming Liu            )
17*3d27d115SPeiming Liu            with InsertionPoint(sequence.body):
18*3d27d115SPeiming Liu                f(sequence.bodyTarget)
19*3d27d115SPeiming Liu                transform.YieldOp()
20*3d27d115SPeiming Liu        print("\nTEST:", f.__name__)
21*3d27d115SPeiming Liu        print(module)
22*3d27d115SPeiming Liu    return f
23*3d27d115SPeiming Liu
24*3d27d115SPeiming Liu
25*3d27d115SPeiming Liu@run
26*3d27d115SPeiming Liudef testMatchSparseInOut(target):
27*3d27d115SPeiming Liu    sparse_tensor.MatchSparseInOut(transform.AnyOpType.get(), target)
28*3d27d115SPeiming Liu    # CHECK-LABEL: TEST: testMatchSparseInOut
29*3d27d115SPeiming Liu    # CHECK:       transform.sequence
30*3d27d115SPeiming Liu    # CHECK-NEXT:  ^{{.*}}(%[[ARG0:.*]]: !transform.any_op):
31*3d27d115SPeiming Liu    # CHECK-NEXT:    transform.sparse_tensor.match.sparse_inout %[[ARG0]]
32