blob: 3e7a8b27d4c09df56823f93f030ab487aa68e81f (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
|
# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
import numpy as np
import mlir.dialects.func as func
import mlir.dialects.shape as shape
def run(f):
print("\nTEST:", f.__name__)
f()
return f
# CHECK-LABEL: TEST: testConstShape
@run
def testConstShape():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
@func.FuncOp.from_py_func(
RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32))
def const_shape_tensor(arg):
shape.ConstWitnessOp(False)
shape.ConstSizeOp(30)
shape.ConstSizeOp(IntegerAttr.get(IndexType.get(), 40))
x = shape.ConstShapeOp([1, 2])
shape.MeetOp(x, x, error="impossible")
return shape.ConstShapeOp(
DenseElementsAttr.get(
np.array([3, 4], dtype=np.int64), type=IndexType.get()))
# CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
# CHECK-DAG: shape.const_witness false
# CHECK-DAG: shape.const_size 30
# CHECK-DAG: shape.const_size 40
# CHECK-DAG: shape.const_shape [1, 2] : tensor<2xindex>
# CHECK-DAG: shape.const_shape [3, 4] : tensor<2xindex>
print(module)
|
