100 lines · plain
1// RUN: mlir-opt --resolve-ranked-shaped-type-result-dims --split-input-file %s | FileCheck %s2 3// CHECK-LABEL: func @dim_out_of_bounds(4// CHECK-NEXT: arith.constant5// CHECK-NEXT: memref.dim6// CHECK-NEXT: return7func.func @dim_out_of_bounds(%m : memref<7x8xf32>) -> index {8 %idx = arith.constant 7 : index9 %0 = memref.dim %m, %idx : memref<7x8xf32>10 return %0 : index11}12 13// -----14 15// CHECK-LABEL: func @dim_out_of_bounds_2(16// CHECK-NEXT: arith.constant17// CHECK-NEXT: arith.constant18// CHECK-NEXT: bufferization.alloc_tensor19// CHECK-NEXT: tensor.dim20// CHECK-NEXT: return21func.func @dim_out_of_bounds_2(%idx1 : index, %idx2 : index) -> index {22 %idx = arith.constant 7 : index23 %sz = arith.constant 5 : index24 %alloc = bufferization.alloc_tensor(%sz, %sz) : tensor<?x?xf32>25 %0 = tensor.dim %alloc, %idx : tensor<?x?xf32>26 return %0 : index27}28 29// -----30 31// CHECK-LABEL: func.func @dynamic_dim_of_transpose_op(32// CHECK-SAME: %[[arg:.*]]: tensor<1x2x?x8xi8>) -> index {33// CHECK-NEXT: %[[c2:.*]] = arith.constant 234// CHECK-NEXT: tensor.dim %[[arg]], %[[c2]]35// CHECK-NEXT: return36func.func @dynamic_dim_of_transpose_op(%arg0: tensor<1x2x?x8xi8>) -> index {37 %1 = tosa.transpose %arg0 { perms = array<i32: 0, 3, 1, 2> }: (tensor<1x2x?x8xi8>) -> tensor<1x8x2x?xi8>38 %c3 = arith.constant 3 : index39 %dim = tensor.dim %1, %c3 : tensor<1x8x2x?xi8>40 return %dim : index41}42 43// -----44 45// CHECK-LABEL: func.func @static_dim_of_transpose_op(46// CHECK: arith.constant 100 : index47// CHECK: return48func.func @static_dim_of_transpose_op(%arg0: tensor<1x100x?x8xi8>) -> index {49 %1 = tosa.transpose %arg0 { perms = array<i32: 0, 3, 1, 2> }: (tensor<1x100x?x8xi8>) -> tensor<1x8x100x?xi8>50 %c2 = arith.constant 2 : index51 %dim = tensor.dim %1, %c2 : tensor<1x8x100x?xi8>52 return %dim : index53}54 55// -----56 57// Test case: Folding of memref.dim(memref.expand_shape)58// CHECK-LABEL: func @dim_of_memref_expand_shape(59// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<?x8xi32>60// CHECK-NEXT: %[[IDX:.*]] = arith.constant 061// CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[MEM]], %[[IDX]] : memref<?x8xi32>62// CHECK: return %[[DIM]] : index63func.func @dim_of_memref_expand_shape(%arg0: memref<?x8xi32>)64 -> index {65 %c0 = arith.constant 0 : index66 %c1 = arith.constant 1 : index67 %s = memref.dim %arg0, %c0 : memref<?x8xi32>68 %0 = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, %s, 2, 4]: memref<?x8xi32> into memref<1x?x2x4xi32>69 %1 = memref.dim %0, %c1 : memref<1x?x2x4xi32>70 return %1 : index71}72 73// -----74 75// CHECK-LABEL: @iter_to_init_arg_loop_like76// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {77// CHECK: %[[RESULT:.*]] = scf.forall78// CHECK-SAME: shared_outs(%[[OUTS:.*]] = %[[ARG1]]) -> (tensor<?x?xf32>) {79// CHECK-NEXT: %{{.*}} = tensor.dim %[[ARG1]], %{{.*}} : tensor<?x?xf32>80func.func @iter_to_init_arg_loop_like(81 %arg0 : tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {82 %c0 = arith.constant 0 : index83 %c1 = arith.constant 1 : index84 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>85 86 %result = scf.forall (%i) = (%c0) to (%dim0)87 step (%c1) shared_outs(%o = %arg1) -> (tensor<?x?xf32>) {88 89 %dim1 = tensor.dim %o, %c1 : tensor<?x?xf32>90 %slice = tensor.extract_slice %arg1[%i, 0] [1, %dim1] [1, 1]91 : tensor<?x?xf32> to tensor<1x?xf32>92 93 scf.forall.in_parallel {94 tensor.parallel_insert_slice %slice into %o[%i, 0] [1, %dim1] [1, 1]95 : tensor<1x?xf32> into tensor<?x?xf32>96 }97 }98 return %result : tensor<?x?xf32>99}100