27 lines
1.3 KiB
MLIR
27 lines
1.3 KiB
MLIR
|
// RUN: onnx-mlir-opt --lower-krnl-shape %s -split-input-file | FileCheck %s
|
||
|
|
|
||
|
|
func @test_krnl_shape_lowering(%arg0: memref<?x?xf32>) -> index {
|
||
|
|
%c1 = constant 1 : index
|
||
|
|
%c0 = constant 0 : index
|
||
|
|
%0 = dim %arg0, %c0 : memref<?x?xf32>
|
||
|
|
%1 = alloc(%0) : memref<?x10xf32>
|
||
|
|
%shape = "krnl.shape"(%1) : (memref<?x10xf32>) -> !shape.shape
|
||
|
|
%d1 = "shape.get_extent"(%shape, %c1) : (!shape.shape, index) -> !shape.size
|
||
|
|
%ind = "shape.size_to_index"(%d1) : (!shape.size) -> index
|
||
|
|
%e = addi %ind, %ind : index
|
||
|
|
return %e : index
|
||
|
|
|
||
|
|
// CHECK-LABEL: test_krnl_shape_lowering
|
||
|
|
// CHECK: [[CONST0:%.+]] = constant 0 : index
|
||
|
|
// CHECK: [[CONST1:%.+]] = constant 1 : index
|
||
|
|
// CHECK: [[DIM:%.+]] = dim %arg0, [[CONST0]] : memref<?x?xf32>
|
||
|
|
// CHECK: [[ALLOC:%.+]] = alloc([[DIM]]) : memref<?x10xf32>
|
||
|
|
// CHECK: [[DIM0:%.+]] = "krnl.dim"([[ALLOC]], [[CONST0]]) : (memref<?x10xf32>, index) -> index
|
||
|
|
// CHECK: [[DIM1:%.+]] = "krnl.dim"([[ALLOC]], [[CONST1]]) : (memref<?x10xf32>, index) -> index
|
||
|
|
// CHECK: [[SHAPE:%.+]] = shape.from_extents [[DIM0]], [[DIM1]]
|
||
|
|
// CHECK: [[EXTENT:%.+]] = shape.get_extent [[SHAPE]], [[CONST1]] : !shape.shape, index -> !shape.size
|
||
|
|
// CHECK: [[EXTENT_AS_INDEX:%.+]] = shape.size_to_index [[EXTENT]] : !shape.size
|
||
|
|
// CHECK: [[RES:%.+]] = addi [[EXTENT_AS_INDEX]], [[EXTENT_AS_INDEX]] : index
|
||
|
|
// CHECK: return [[RES]] : index
|
||
|
|
}
|