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