xref: /llvm-project/mlir/test/python/dialects/transform_tensor_ext.py (revision 92233062c17590d3157bdc6db430fcdfc54312fe) 
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import transform
5from mlir.dialects.transform import tensor
6
7
8def run(f):
9    print("\nTEST:", f.__name__)
10    with Context(), Location.unknown():
11        module = Module.create()
12        with InsertionPoint(module.body):
13            sequence = transform.SequenceOp(
14                transform.FailurePropagationMode.Propagate,
15                [],
16                transform.AnyOpType.get(),
17            )
18            with InsertionPoint(sequence.body):
19                f(sequence.bodyTarget)
20                transform.YieldOp()
21        print(module)
22    return f
23
24
25@run
26def testMakeLoopIndependentOpCompact(target):
27    tensor.MakeLoopIndependentOp(target, 4)
28    # CHECK-LABEL: TEST: testMakeLoopIndependentOpCompact
29    # CHECK: = transform.tensor.make_loop_independent
30    # CHECK-SAME: num_loops = 4 : i64
31    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
32
33
34@run
35def testMakeLoopIndependentOpTyped(target):
36    tensor.MakeLoopIndependentOp(transform.OperationType.get("test.dummy"), target, 4)
37    # CHECK-LABEL: TEST: testMakeLoopIndependentOpTyped
38    # CHECK: = transform.tensor.make_loop_independent
39    # CHECK-SAME: num_loops = 4 : i64
40    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
41