154 lines · plain
1// RUN: mlir-opt --transform-interpreter -canonicalize -split-input-file --verify-diagnostics %s | FileCheck %s2 3// CHECK-LABEL: pad_fill4// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : tensor<8x25xf32>) -> tensor<8x25xf32>5func.func @pad_fill(%value: f32, %output: tensor<24x25xf32>) -> tensor<24x25xf32>6{7 %0 = linalg.fill ins(%value : f32) outs(%output : tensor<24x25xf32>) -> tensor<24x25xf32>8 func.return %0 : tensor<24x25xf32>9}10 11module attributes {transform.with_named_sequence} {12 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {13 %fill = transform.structured.match ops{["linalg.fill"]} in %arg114 : (!transform.any_op) -> !transform.any_op15 16 // Tile to 5 then pad to 817 %fill_l1, %loops_l1 = transform.structured.tile_using_for %fill tile_sizes [5]18 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)19 20 %fill_padded, %_ = transform.structured.pad_tiling_interface %fill_l1 to padding_sizes [8] {21 padding_values= [#ub.poison, 0.0 : f32]22 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)23 24 transform.yield25 }26}27 28// -----29 30// CHECK-LABEL: pad_lhs31func.func @pad_lhs(32 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)33 -> tensor<24x25xf32>34{35 // CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>)36 // CHECK: tensor.pad %{{.*}}37 // CHECK: : tensor<?x12xf32> to tensor<8x12xf32>38 // CHECK: tensor.pad %{{.*}}39 // CHECK: : tensor<?x25xf32> to tensor<8x25xf32>40 // CHECK: linalg.matmul ins(%{{.*}}, %{{.*}} : tensor<8x12xf32>, tensor<12x25xf32>) outs(%{{.*}} : tensor<8x25xf32>) -> tensor<8x25xf32>41 // CHECK: tensor.extract_slice %{{.*}}[0, 0] [%{{.*}}, 25] [1, 1]42 // CHECK: : tensor<8x25xf32> to tensor<?x25xf32>43 // CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0] [%{{.*}}, 25] [1, 1]44 // CHECK-SAME: : tensor<?x25xf32> into tensor<24x25xf32>45 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>46 func.return %0 : tensor<24x25xf32>47}48 49module attributes {transform.with_named_sequence} {50 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {51 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg152 : (!transform.any_op) -> !transform.any_op53 54 // Tile to 5 then pad to 8 (supposedly to better hit vector ops).55 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)56 %matmul_padded, %_ = transform.structured.pad_tiling_interface %matmul_l1 to padding_sizes [8] {57 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]58 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)59 60 transform.yield61 }62}63 64// -----65 66#map = affine_map<(d0, d1, d2) -> (d0, d1)>67#map1 = affine_map<(d0, d1, d2) -> (d0, d2, d0 + d1)>68module {69 70// CHECK-LABEL: @generic71// CHECK-SAME: %[[T0:.*]]: tensor<7x5xf32>,72// CHECK-SAME: %[[T1:.*]]: tensor<7x11x11xf32>)73 func.func @generic(%arg0: tensor<7x5xf32>, %arg1: tensor<7x11x11xf32>) -> tensor<7x11x11xf32> {74 75 // CHECK-DAG: %[[CST:.*]] = arith.constant 0.76 77 // CHECK: %[[PAD0:.*]] = tensor.pad %[[T0]] low[0, 0] high[1, 0]78 // CHECK: : tensor<7x5xf32> to tensor<8x5xf32>79 // CHECK: %[[PAD1:.*]] = tensor.pad %[[T1]] low[0, 0, 0] high[1, 3, 1] {80 // CHECK: : tensor<7x11x11xf32> to tensor<8x14x12xf32>81 // CHECK-NEXT: linalg.generic82 // CHECK: tensor.extract_slice %{{.*}}[0, 0, 0] [7, 11, 11] [1, 1, 1] : tensor<8x14x12xf32> to tensor<7x11x11xf32>83 %0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<7x5xf32>) outs(%arg1 : tensor<7x11x11xf32>) {84 ^bb0(%in: f32, %out: f32):85 linalg.yield %in : f3286 } -> tensor<7x11x11xf32>87 return %0 : tensor<7x11x11xf32>88 }89 module attributes {transform.with_named_sequence} {90 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {91 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op92 %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 0, 14] {93 padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]94 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)95 transform.yield96 }97 }98}99 100 101// -----102 103 104// CHECK-DAG: #[[$MAP0:.*]] = affine_map<()[s0] -> (-s0 + 8)>105// CHECK-DAG: #[[$MAP1:.*]] = affine_map<()[s0] -> (-s0 + 12)>106// CHECK-DAG: #[[$MAP2:.*]] = affine_map<()[s0] -> (s0 + 5)>107 108#map = affine_map<(d0, d1, d2) -> (d0, d1)>109#map1 = affine_map<(d0, d1, d2) -> (d0, d2, d0 + d1)>110module {111 112// CHECK-LABEL: @generic113// CHECK-SAME: %[[T0:.*]]: tensor<?x5xf32>,114// CHECK-SAME: %[[T1:.*]]: tensor<?x11x?xf32>)115 func.func @generic(%arg0: tensor<?x5xf32>, %arg1: tensor<?x11x?xf32>) -> tensor<?x11x?xf32> {116 117 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index118 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index119 //120 // CHECK: %[[D0_0:.*]] = tensor.dim %{{.*}}, %[[C0]] : tensor<?x5xf32>121 // CHECK: %[[H0:.*]] = affine.apply #[[$MAP0]]()[%[[D0_0]]]122 // CHECK: tensor.pad %{{.*}} low[0, 0] high[%[[H0]], 0] {123 // CHECK: : tensor<?x5xf32> to tensor<8x5xf32>124 //125 // CHECK: %[[D0_1:.*]] = tensor.dim %{{.*}}, %[[C0]] : tensor<?x11x?xf32>126 // CHECK: %[[H1:.*]] = affine.apply #[[$MAP0]]()[%[[D0_1]]]127 // CHECK: %[[D2_0:.*]] = tensor.dim %{{.*}}, %[[C2]] : tensor<?x11x?xf32>128 // CHECK: %[[H2:.*]] = affine.apply #[[$MAP1]]()[%[[D2_0]]]129 // CHECK: tensor.pad %{{.*}} low[0, 0, 0] high[%[[H1]], 3, %[[H2]]] {130 // CHECK: : tensor<?x11x?xf32> to tensor<8x14x12xf32>131 //132 // CHECK: %[[D0_2:.*]] = tensor.dim %{{.*}}, %[[C0]] : tensor<?x5xf32>133 // CHECK: %[[D2_1:.*]] = affine.apply #[[$MAP2]]()[%[[D0_2]]]134 // CHECK: linalg.generic {{.*}} ins(%{{.*}} : tensor<8x5xf32>) outs(%{{.*}} : tensor<8x14x12xf32>) {135 // CHECK: } -> tensor<8x14x12xf32>136 // CHECK: tensor.extract_slice %{{.*}}[0, 0, 0] [%[[D0_2]], 11, %[[D2_1]]] [1, 1, 1] : tensor<8x14x12xf32> to tensor<?x11x?xf32>137 //138 %0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x5xf32>) outs(%arg1 : tensor<?x11x?xf32>) {139 ^bb0(%in: f32, %out: f32):140 linalg.yield %in : f32141 } -> tensor<?x11x?xf32>142 return %0 : tensor<?x11x?xf32>143 }144 module attributes {transform.with_named_sequence} {145 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {146 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op147 %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 0, 14] {148 padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]149 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)150 transform.yield151 }152 }153}154