23aa5a7446
The last remaining operations in the standard dialect all revolve around FuncOp/function related constructs. This patch simply handles the initial renaming (which by itself is already huge), but there are a large number of cleanups unlocked/necessary afterwards: * Removing a bunch of unnecessary dependencies on Func * Cleaning up the From/ToStandard conversion passes * Preparing for the move of FuncOp to the Func dialect See the discussion at https://discourse.llvm.org/t/standard-dialect-the-final-chapter/6061 Differential Revision: https://reviews.llvm.org/D120624
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import arith
|
|
from mlir.dialects import func
|
|
from mlir.dialects import scf
|
|
from mlir.dialects import builtin
|
|
|
|
|
|
def constructAndPrintInModule(f):
|
|
print("\nTEST:", f.__name__)
|
|
with Context(), Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
f()
|
|
print(module)
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testSimpleLoop
|
|
@constructAndPrintInModule
|
|
def testSimpleLoop():
|
|
index_type = IndexType.get()
|
|
|
|
@builtin.FuncOp.from_py_func(index_type, index_type, index_type)
|
|
def simple_loop(lb, ub, step):
|
|
loop = scf.ForOp(lb, ub, step, [lb, lb])
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp(loop.inner_iter_args)
|
|
return
|
|
|
|
|
|
# CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
|
|
# CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
|
|
# CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
|
|
# CHECK: scf.yield %[[I1]], %[[I2]]
|
|
|
|
|
|
# CHECK-LABEL: TEST: testInductionVar
|
|
@constructAndPrintInModule
|
|
def testInductionVar():
|
|
index_type = IndexType.get()
|
|
|
|
@builtin.FuncOp.from_py_func(index_type, index_type, index_type)
|
|
def induction_var(lb, ub, step):
|
|
loop = scf.ForOp(lb, ub, step, [lb])
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp([loop.induction_variable])
|
|
return
|
|
|
|
|
|
# CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
|
|
# CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
|
|
# CHECK: scf.yield %[[IV]]
|
|
|
|
|
|
@constructAndPrintInModule
|
|
def testOpsAsArguments():
|
|
index_type = IndexType.get()
|
|
callee = builtin.FuncOp(
|
|
"callee", ([], [index_type, index_type]), visibility="private")
|
|
f = builtin.FuncOp("ops_as_arguments", ([], []))
|
|
with InsertionPoint(f.add_entry_block()):
|
|
lb = arith.ConstantOp.create_index(0)
|
|
ub = arith.ConstantOp.create_index(42)
|
|
step = arith.ConstantOp.create_index(2)
|
|
iter_args = func.CallOp(callee, [])
|
|
loop = scf.ForOp(lb, ub, step, iter_args)
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp(loop.inner_iter_args)
|
|
func.ReturnOp([])
|
|
|
|
|
|
# CHECK-LABEL: TEST: testOpsAsArguments
|
|
# CHECK: func private @callee() -> (index, index)
|
|
# CHECK: func @ops_as_arguments() {
|
|
# CHECK: %[[LB:.*]] = arith.constant 0
|
|
# CHECK: %[[UB:.*]] = arith.constant 42
|
|
# CHECK: %[[STEP:.*]] = arith.constant 2
|
|
# CHECK: %[[ARGS:.*]]:2 = call @callee()
|
|
# CHECK: scf.for %arg0 = %c0 to %c42 step %c2
|
|
# CHECK: iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
|
|
# CHECK: scf.yield %{{.*}}, %{{.*}}
|
|
# CHECK: return
|