32 lines · plain
1// RUN: mlir-opt -reify-result-shapes %s | FileCheck %s2 3// The test below checks concat op reification. In the first case, no cast is inserted while on the second a cast gets inserted.4// CHECK-LABEL: func.func @concat_reification5func.func @concat_reification(%arg0: tensor<4x7x3xf32>, %arg1 : tensor<4x4x3xf32>, %arg2: tensor<?x?x?xf32>)6 -> (tensor<4x11x3xf32>, tensor<?x?x?xf32>) {7 // CHECK: %[[RES0:.*]] = tensor.concat dim(1) %{{.*}} : (tensor<4x7x3xf32>, tensor<4x4x3xf32>) -> tensor<4x11x3xf32>8 %1 = tensor.concat dim(1) %arg0, %arg1 : (tensor<4x7x3xf32>, tensor<4x4x3xf32>) -> tensor<4x11x3xf32>9 // CHECK: %[[V0:.*]] = tensor.concat dim(2) %{{.*}} : (tensor<4x7x3xf32>, tensor<?x?x?xf32>) -> tensor<4x7x?xf32>10 // CHECK: %[[RES1:.*]] = tensor.cast %[[V0]] : tensor<4x7x?xf32> to tensor<?x?x?xf32>11 %2 = tensor.concat dim(2) %arg0, %arg2 : (tensor<4x7x3xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>12 // CHECK: return %[[RES0]], %[[RES1]] : tensor<4x11x3xf32>, tensor<?x?x?xf32>13 return %1, %2 : tensor<4x11x3xf32>, tensor<?x?x?xf32>14}15 16// CHECK-LABEL: func.func @pad_reification17func.func @pad_reification(%cst : f32, %idx : index, %t: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {18 %pad_amt = affine.apply affine_map<(d0) -> (-d0 + 256)>(%idx)19 %es = tensor.extract_slice %t[0, 0, 0] [1, %idx, 64] [1, 1, 1] 20 : tensor<64x?x64xf32> to tensor<1x?x64xf32>21 22 // CHECK: tensor.pad23 // CHECK: : tensor<1x?x64xf32> to tensor<1x256x64xf32>24 // CHECK: tensor.cast %{{.*}} : tensor<1x256x64xf32> to tensor<1x?x64xf32>25 %padded = tensor.pad %es low[0, 0, 0] high[0, %pad_amt, 0] {26 ^bb0(%a: index, %b: index, %c: index):27 tensor.yield %cst : f3228 } : tensor<1x?x64xf32> to tensor<1x?x64xf32>29 30 return %padded : tensor<1x?x64xf32>31}32