129 lines · plain
1// RUN: mlir-opt -split-input-file -transform-interpreter %s | FileCheck %s2 3module attributes {transform.with_named_sequence} {4 transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {5 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">6 transform.apply_patterns to %func_op {7 transform.apply_patterns.tensor.fold_tensor_empty8 } : !transform.op<"func.func">9 transform.yield10 }11}12 13// CHECK: #[[$MAP2:.+]] = affine_map<()[s0] -> (s0 * 28)>14 15func.func @empty_reshape_expansion(%arg0 : index, %sz0: index) -> tensor<2x3x5x4x?x7xf32> {16 %0 = tensor.empty(%arg0) : tensor<6x5x?xf32>17 %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]] output_shape [2, 3, 5, 4, %sz0, 7] : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>18 return %1 : tensor<2x3x5x4x?x7xf32>19}20// CHECK-LABEL: func @empty_reshape_expansion21// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index,22// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index23// CHECK-NEXT: %[[INIT:.+]] = tensor.empty(%[[ARG1]])24// CHECK-NEXT: return %[[INIT]]25 26func.func @empty_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {27 %0 = tensor.empty(%arg0) : tensor<2x3x5x4x?x7xf32>28 %1 = tensor.collapse_shape %0 [[0, 1], [2], [3, 4, 5]]29 : tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>30 return %1 : tensor<6x5x?xf32>31}32// CHECK-LABEL: func @empty_reshape_collapse33// CHECK-SAME: %[[ARG0:.+]]: index34// CHECK: %[[OLD_INIT:.+]] = tensor.empty(%{{.*}}) : tensor<2x3x5x4x?x7xf32>35// CHECK-NEXT: %[[DIM:.*]] = tensor.dim %[[OLD_INIT]]36// CHECK-NEXT: %[[D:.+]] = affine.apply #[[$MAP2]]()[%[[DIM]]]37// CHECK-NEXT: %[[INIT:.+]] = tensor.empty(%[[D]])38// CHECK-NEXT: return %[[INIT]]39 40func.func @fold_empty_tensor_with_slice41 (%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>42{43 %0 = tensor.empty(%arg0) : tensor<?x10x40xf32>44 %1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]45 : tensor<?x10x40xf32> to tensor<5x?x20xf32>46 return %1 : tensor<5x?x20xf32>47}48// CHECK-LABEL: func @fold_empty_tensor_with_slice49// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index50// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index51// CHECK: %[[T0:.+]] = tensor.empty(%[[ARG1]])52// CHECK: return %[[T0]]53 54// CHECK-LABEL: func @rank_reducing_empty_tensor_extract55func.func @rank_reducing_empty_tensor_extract(%sz : index, %idx : index) -> tensor<2xf32> {56 // CHECK: tensor.empty() : tensor<2xf32>57 %a = tensor.empty(%sz) : tensor<?x2xf32>58 59 // CHECK-NOT: extract60 %r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>61 return %r: tensor<2xf32>62}63 64// -----65 66module attributes {transform.with_named_sequence} {67 transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {68 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">69 transform.apply_patterns to %func_op {70 transform.apply_patterns.tensor.fold_tensor_empty71 {fold_single_use_only = true}72 } : !transform.op<"func.func">73 transform.yield74 }75}76 77func.func @double_use_of_tensor_empty(%arg0: index, %arg1: index)78 -> (tensor<5x?x20xf32>, tensor<5x?x20xf32>)79{80 %0 = tensor.empty(%arg0) : tensor<?x10x40xf32>81 %1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]82 : tensor<?x10x40xf32> to tensor<5x?x20xf32>83 %2 = tensor.extract_slice %0[1, 1, 1] [5, %arg1, 20] [1, 1, 1]84 : tensor<?x10x40xf32> to tensor<5x?x20xf32>85 return %1, %2 : tensor<5x?x20xf32>, tensor<5x?x20xf32>86}87// CHECK-LABEL: func @double_use_of_tensor_empty(88// CHECK: tensor.empty{{.*}} : tensor<?x10x40xf32>89// CHECK: tensor.extract_slice90// CHECK: tensor.extract_slice91 92// -----93 94module attributes {transform.with_named_sequence} {95 transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {96 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">97 transform.apply_patterns to %func_op {98 transform.apply_patterns.tensor.fold_tensor_empty99 } : !transform.op<"func.func">100 transform.yield101 }102}103 104func.func @concats_of_empty(105 %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index)106 -> tensor<5x?x?xf32>107{108 %0 = tensor.empty(%arg0, %arg1) : tensor<5x?x?xf32>109 %1 = tensor.empty(%arg2, %arg3) : tensor<5x?x?xf32>110 %2 = tensor.concat dim(1) %0, %1 : (tensor<5x?x?xf32>, tensor<5x?x?xf32>) -> tensor<5x?x?xf32>111 return %2 : tensor<5x?x?xf32>112}113// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>114// CHECK: func @concats_of_empty(115// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index,116// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index,117// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index,118// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index)119// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index120// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index121// CHECK-DAG: %[[EMPTY0:.+]] = tensor.empty(%[[ARG0]], %[[ARG1]])122// CHECK-DAG: %[[EMPTY1:.+]] = tensor.empty(%[[ARG2]], %[[ARG3]])123// CHECK: %[[D2:.+]] = tensor.dim %[[EMPTY0]], %[[C2]]124// CHECK-DAG: %[[D0_1:.+]] = tensor.dim %[[EMPTY0]], %[[C1]]125// CHECK-DAG: %[[D1_1:.+]] = tensor.dim %[[EMPTY1]], %[[C1]]126// CHECK-DAG: %[[SUM:.+]] = affine.apply #[[MAP]]()[%[[D0_1]], %[[D1_1]]]127// CHECK: %[[NEW_EMPTY:.+]] = tensor.empty(%[[SUM]], %[[D2]])128// CHECK: return %[[NEW_EMPTY]]129