xref: /llvm-project/mlir/test/python/dialects/transform_tensor_ext.py (revision 92233062c17590d3157bdc6db430fcdfc54312fe)
10575ab2dSIngo Müller# RUN: %PYTHON %s | FileCheck %s
20575ab2dSIngo Müller
30575ab2dSIngo Müllerfrom mlir.ir import *
40575ab2dSIngo Müllerfrom mlir.dialects import transform
50575ab2dSIngo Müllerfrom mlir.dialects.transform import tensor
60575ab2dSIngo Müller
70575ab2dSIngo Müller
80575ab2dSIngo Müllerdef run(f):
90575ab2dSIngo Müller    print("\nTEST:", f.__name__)
100575ab2dSIngo Müller    with Context(), Location.unknown():
110575ab2dSIngo Müller        module = Module.create()
120575ab2dSIngo Müller        with InsertionPoint(module.body):
130575ab2dSIngo Müller            sequence = transform.SequenceOp(
14*92233062Smax                transform.FailurePropagationMode.Propagate,
150575ab2dSIngo Müller                [],
160575ab2dSIngo Müller                transform.AnyOpType.get(),
170575ab2dSIngo Müller            )
180575ab2dSIngo Müller            with InsertionPoint(sequence.body):
190575ab2dSIngo Müller                f(sequence.bodyTarget)
200575ab2dSIngo Müller                transform.YieldOp()
210575ab2dSIngo Müller        print(module)
220575ab2dSIngo Müller    return f
230575ab2dSIngo Müller
240575ab2dSIngo Müller
250575ab2dSIngo Müller@run
260575ab2dSIngo Müllerdef testMakeLoopIndependentOpCompact(target):
270575ab2dSIngo Müller    tensor.MakeLoopIndependentOp(target, 4)
280575ab2dSIngo Müller    # CHECK-LABEL: TEST: testMakeLoopIndependentOpCompact
290575ab2dSIngo Müller    # CHECK: = transform.tensor.make_loop_independent
300575ab2dSIngo Müller    # CHECK-SAME: num_loops = 4 : i64
310575ab2dSIngo Müller    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
320575ab2dSIngo Müller
330575ab2dSIngo Müller
340575ab2dSIngo Müller@run
350575ab2dSIngo Müllerdef testMakeLoopIndependentOpTyped(target):
360575ab2dSIngo Müller    tensor.MakeLoopIndependentOp(transform.OperationType.get("test.dummy"), target, 4)
370575ab2dSIngo Müller    # CHECK-LABEL: TEST: testMakeLoopIndependentOpTyped
380575ab2dSIngo Müller    # CHECK: = transform.tensor.make_loop_independent
390575ab2dSIngo Müller    # CHECK-SAME: num_loops = 4 : i64
400575ab2dSIngo Müller    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
41