922b26cde4
This includes some minor customization for FuncOp and ModuleOp. Differential Revision: https://reviews.llvm.org/D95022
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
import mlir.dialects.builtin as builtin
|
|
import mlir.dialects.std as std
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
|
|
|
|
# CHECK-LABEL: TEST: testBuildFuncOp
|
|
def testBuildFuncOp():
|
|
ctx = Context()
|
|
with Location.unknown(ctx) as loc:
|
|
m = builtin.ModuleOp()
|
|
|
|
f32 = F32Type.get()
|
|
tensor_type = RankedTensorType.get((2, 3, 4), f32)
|
|
with InsertionPoint.at_block_begin(m.body):
|
|
func = builtin.FuncOp(name="some_func",
|
|
type=FunctionType.get(
|
|
inputs=[tensor_type, tensor_type],
|
|
results=[tensor_type]),
|
|
visibility="nested")
|
|
# CHECK: Name is: "some_func"
|
|
print("Name is: ", func.name)
|
|
|
|
# CHECK: Type is: (tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
|
|
print("Type is: ", func.type)
|
|
|
|
# CHECK: Visibility is: "nested"
|
|
print("Visibility is: ", func.visibility)
|
|
|
|
try:
|
|
entry_block = func.entry_block
|
|
except IndexError as e:
|
|
# CHECK: External function does not have a body
|
|
print(e)
|
|
|
|
with InsertionPoint(func.add_entry_block()):
|
|
std.ReturnOp([func.entry_block.arguments[0]])
|
|
pass
|
|
|
|
try:
|
|
func.add_entry_block()
|
|
except IndexError as e:
|
|
# CHECK: The function already has an entry block!
|
|
print(e)
|
|
|
|
# Try the callback builder and passing type as tuple.
|
|
func = builtin.FuncOp(name="some_other_func",
|
|
type=([tensor_type, tensor_type], [tensor_type]),
|
|
visibility="nested",
|
|
body_builder=lambda func: std.ReturnOp(
|
|
[func.entry_block.arguments[0]]))
|
|
|
|
# CHECK: module {
|
|
# CHECK: func nested @some_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
|
|
# CHECK: return %arg0 : tensor<2x3x4xf32>
|
|
# CHECK: }
|
|
# CHECK: func nested @some_other_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
|
|
# CHECK: return %arg0 : tensor<2x3x4xf32>
|
|
# CHECK: }
|
|
print(m)
|
|
|
|
|
|
run(testBuildFuncOp)
|