252 lines
10 KiB
Python
252 lines
10 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import mlir.dialects.arith as arith
|
|
import mlir.dialects.memref as memref
|
|
import mlir.extras.types as T
|
|
from mlir.dialects.memref import _infer_memref_subview_result_type
|
|
from mlir.ir import *
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testSubViewAccessors
|
|
@run
|
|
def testSubViewAccessors():
|
|
ctx = Context()
|
|
module = Module.parse(
|
|
r"""
|
|
func.func @f1(%arg0: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = arith.constant 1 : index
|
|
%2 = arith.constant 2 : index
|
|
%3 = arith.constant 3 : index
|
|
%4 = arith.constant 4 : index
|
|
%5 = arith.constant 5 : index
|
|
memref.subview %arg0[%0, %1][%2, %3][%4, %5] : memref<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
|
|
return
|
|
}
|
|
""",
|
|
ctx,
|
|
)
|
|
func_body = module.body.operations[0].regions[0].blocks[0]
|
|
subview = func_body.operations[6]
|
|
|
|
assert subview.source == subview.operands[0]
|
|
assert len(subview.offsets) == 2
|
|
assert len(subview.sizes) == 2
|
|
assert len(subview.strides) == 2
|
|
assert subview.result == subview.results[0]
|
|
|
|
# CHECK: SubViewOp
|
|
print(type(subview).__name__)
|
|
|
|
# CHECK: constant 0
|
|
print(subview.offsets[0])
|
|
# CHECK: constant 1
|
|
print(subview.offsets[1])
|
|
# CHECK: constant 2
|
|
print(subview.sizes[0])
|
|
# CHECK: constant 3
|
|
print(subview.sizes[1])
|
|
# CHECK: constant 4
|
|
print(subview.strides[0])
|
|
# CHECK: constant 5
|
|
print(subview.strides[1])
|
|
|
|
|
|
# CHECK-LABEL: TEST: testCustomBuidlers
|
|
@run
|
|
def testCustomBuidlers():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
module = Module.parse(
|
|
r"""
|
|
func.func @f1(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index) {
|
|
return
|
|
}
|
|
"""
|
|
)
|
|
f = module.body.operations[0]
|
|
func_body = f.regions[0].blocks[0]
|
|
with InsertionPoint.at_block_terminator(func_body):
|
|
memref.LoadOp(f.arguments[0], f.arguments[1:])
|
|
|
|
# CHECK: func @f1(%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
|
|
# CHECK: memref.load %[[ARG0]][%[[ARG1]], %[[ARG2]]]
|
|
print(module)
|
|
assert module.operation.verify()
|
|
|
|
|
|
# CHECK-LABEL: TEST: testMemRefAttr
|
|
@run
|
|
def testMemRefAttr():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
memref.global_("objFifo_in0", T.memref(16, T.i32()))
|
|
# CHECK: memref.global @objFifo_in0 : memref<16xi32>
|
|
print(module)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testSubViewOpInferReturnTypeSemantics
|
|
@run
|
|
def testSubViewOpInferReturnTypeSemantics():
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
x = memref.alloc(T.memref(10, 10, T.i32()), [], [])
|
|
# CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<10x10xi32>
|
|
print(x.owner)
|
|
|
|
y = memref.subview(x, [1, 1], [3, 3], [1, 1])
|
|
assert y.owner.verify()
|
|
# CHECK: %{{.*}} = memref.subview %[[ALLOC]][1, 1] [3, 3] [1, 1] : memref<10x10xi32> to memref<3x3xi32, strided<[10, 1], offset: 11>>
|
|
print(y.owner)
|
|
|
|
z = memref.subview(
|
|
x,
|
|
[arith.constant(T.index(), 1), 1],
|
|
[3, 3],
|
|
[1, 1],
|
|
)
|
|
# CHECK: %{{.*}} = memref.subview %[[ALLOC]][1, 1] [3, 3] [1, 1] : memref<10x10xi32> to memref<3x3xi32, strided<[10, 1], offset: 11>>
|
|
print(z.owner)
|
|
|
|
z = memref.subview(
|
|
x,
|
|
[arith.constant(T.index(), 3), arith.constant(T.index(), 4)],
|
|
[3, 3],
|
|
[1, 1],
|
|
)
|
|
# CHECK: %{{.*}} = memref.subview %[[ALLOC]][3, 4] [3, 3] [1, 1] : memref<10x10xi32> to memref<3x3xi32, strided<[10, 1], offset: 34>>
|
|
print(z.owner)
|
|
|
|
s = arith.addi(arith.constant(T.index(), 3), arith.constant(T.index(), 4))
|
|
z = memref.subview(
|
|
x,
|
|
[s, 0],
|
|
[3, 3],
|
|
[1, 1],
|
|
)
|
|
# CHECK: {{.*}} = memref.subview %[[ALLOC]][%0, 0] [3, 3] [1, 1] : memref<10x10xi32> to memref<3x3xi32, strided<[10, 1], offset: ?>>
|
|
print(z)
|
|
|
|
try:
|
|
_infer_memref_subview_result_type(
|
|
x.type,
|
|
[arith.constant(T.index(), 3), arith.constant(T.index(), 4)],
|
|
[ShapedType.get_dynamic_size(), 3],
|
|
[1, 1],
|
|
)
|
|
except ValueError as e:
|
|
# CHECK: Only inferring from python or mlir integer constant is supported
|
|
print(e)
|
|
|
|
try:
|
|
memref.subview(
|
|
x,
|
|
[arith.constant(T.index(), 3), arith.constant(T.index(), 4)],
|
|
[ShapedType.get_dynamic_size(), 3],
|
|
[1, 1],
|
|
)
|
|
except ValueError as e:
|
|
# CHECK: mixed static/dynamic offset/sizes/strides requires explicit result type
|
|
print(e)
|
|
|
|
layout = StridedLayoutAttr.get(ShapedType.get_dynamic_size(), [10, 1])
|
|
x = memref.alloc(
|
|
T.memref(
|
|
10,
|
|
10,
|
|
T.i32(),
|
|
layout=layout,
|
|
),
|
|
[],
|
|
[arith.constant(T.index(), 42)],
|
|
)
|
|
# CHECK: %[[DYNAMICALLOC:.*]] = memref.alloc()[%c42] : memref<10x10xi32, strided<[10, 1], offset: ?>>
|
|
print(x.owner)
|
|
y = memref.subview(
|
|
x,
|
|
[1, 1],
|
|
[3, 3],
|
|
[1, 1],
|
|
result_type=T.memref(3, 3, T.i32(), layout=layout),
|
|
)
|
|
# CHECK: %{{.*}} = memref.subview %[[DYNAMICALLOC]][1, 1] [3, 3] [1, 1] : memref<10x10xi32, strided<[10, 1], offset: ?>> to memref<3x3xi32, strided<[10, 1], offset: ?>>
|
|
print(y.owner)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testSubViewOpInferReturnTypeExtensiveSlicing
|
|
@run
|
|
def testSubViewOpInferReturnTypeExtensiveSlicing():
|
|
def check_strides_offset(memref, np_view):
|
|
layout = memref.type.layout
|
|
dtype_size_in_bytes = np_view.dtype.itemsize
|
|
golden_strides = (np.array(np_view.strides) // dtype_size_in_bytes).tolist()
|
|
golden_offset = (
|
|
np_view.ctypes.data - np_view.base.ctypes.data
|
|
) // dtype_size_in_bytes
|
|
|
|
assert (layout.strides, layout.offset) == (golden_strides, golden_offset)
|
|
|
|
with Context() as ctx, Location.unknown(ctx):
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
shape = (10, 22, 3, 44)
|
|
golden_mem = np.zeros(shape, dtype=np.int32)
|
|
mem1 = memref.alloc(T.memref(*shape, T.i32()), [], [])
|
|
|
|
# fmt: off
|
|
check_strides_offset(memref.subview(mem1, (1, 0, 0, 0), (1, 22, 3, 44), (1, 1, 1, 1)), golden_mem[1:2, ...])
|
|
check_strides_offset(memref.subview(mem1, (0, 1, 0, 0), (10, 1, 3, 44), (1, 1, 1, 1)), golden_mem[:, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (0, 0, 1, 0), (10, 22, 1, 44), (1, 1, 1, 1)), golden_mem[:, :, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (0, 0, 0, 1), (10, 22, 3, 1), (1, 1, 1, 1)), golden_mem[:, :, :, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (0, 1, 0, 1), (10, 1, 3, 1), (1, 1, 1, 1)), golden_mem[:, 1:2, :, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (1, 0, 0, 1), (1, 22, 3, 1), (1, 1, 1, 1)), golden_mem[1:2, :, :, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (1, 1, 0, 0), (1, 1, 3, 44), (1, 1, 1, 1)), golden_mem[1:2, 1:2, :, :])
|
|
check_strides_offset(memref.subview(mem1, (0, 0, 1, 1), (10, 22, 1, 1), (1, 1, 1, 1)), golden_mem[:, :, 1:2, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (0, 1, 1, 0), (10, 1, 1, 44), (1, 1, 1, 1)), golden_mem[:, 1:2, 1:2, :])
|
|
check_strides_offset(memref.subview(mem1, (1, 0, 1, 0), (1, 22, 1, 44), (1, 1, 1, 1)), golden_mem[1:2, :, 1:2, :])
|
|
check_strides_offset(memref.subview(mem1, (1, 1, 0, 1), (1, 1, 3, 1), (1, 1, 1, 1)), golden_mem[1:2, 1:2, :, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (1, 0, 1, 1), (1, 22, 1, 1), (1, 1, 1, 1)), golden_mem[1:2, :, 1:2, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (0, 1, 1, 1), (10, 1, 1, 1), (1, 1, 1, 1)), golden_mem[:, 1:2, 1:2, 1:2])
|
|
check_strides_offset(memref.subview(mem1, (1, 1, 1, 0), (1, 1, 1, 44), (1, 1, 1, 1)), golden_mem[1:2, 1:2, 1:2, :])
|
|
# fmt: on
|
|
|
|
# default strides and offset means no stridedlayout attribute means affinemap layout
|
|
assert memref.subview(
|
|
mem1, (0, 0, 0, 0), (10, 22, 3, 44), (1, 1, 1, 1)
|
|
).type.layout == AffineMapAttr.get(
|
|
AffineMap.get(
|
|
4,
|
|
0,
|
|
[
|
|
AffineDimExpr.get(0),
|
|
AffineDimExpr.get(1),
|
|
AffineDimExpr.get(2),
|
|
AffineDimExpr.get(3),
|
|
],
|
|
)
|
|
)
|
|
|
|
shape = (7, 22, 30, 44)
|
|
golden_mem = np.zeros(shape, dtype=np.int32)
|
|
mem2 = memref.alloc(T.memref(*shape, T.i32()), [], [])
|
|
# fmt: off
|
|
check_strides_offset(memref.subview(mem2, (0, 0, 0, 0), (7, 11, 3, 44), (1, 2, 1, 1)), golden_mem[:, 0:22:2])
|
|
check_strides_offset(memref.subview(mem2, (0, 0, 0, 0), (7, 11, 11, 44), (1, 2, 30, 1)), golden_mem[:, 0:22:2, 0:330:30])
|
|
check_strides_offset(memref.subview(mem2, (0, 0, 0, 0), (7, 11, 11, 11), (1, 2, 30, 400)), golden_mem[:, 0:22:2, 0:330:30, 0:4400:400])
|
|
# fmt: on
|
|
|
|
shape = (8, 8)
|
|
golden_mem = np.zeros(shape, dtype=np.int32)
|
|
# fmt: off
|
|
mem3 = memref.alloc(T.memref(*shape, T.i32()), [], [])
|
|
check_strides_offset(memref.subview(mem3, (0, 0), (4, 4), (1, 1)), golden_mem[0:4, 0:4])
|
|
check_strides_offset(memref.subview(mem3, (4, 4), (4, 4), (1, 1)), golden_mem[4:8, 4:8])
|
|
# fmt: on
|