212 lines · plain
1// RUN: mlir-opt --transform-interpreter -canonicalize -split-input-file --verify-diagnostics %s | FileCheck %s2 3func.func @pad_and_hoist_rhs(4 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)5 -> tensor<24x25xf32>6{7 // expected-note @below {{payload operation}}8 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>9 func.return %0 : tensor<24x25xf32>10}11 12module attributes {transform.with_named_sequence} {13 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {14 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg115 : (!transform.any_op) -> !transform.any_op16 17 18 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)19 20 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {21 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],22 padding_dimensions=[0, 1, 2],23 copy_back_op = "none"24 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)25 26 // In this case, the pad op is actually empty: we only tile the first dimension27 // and it does not have an impact on the RHS operand.28 // expected-error @below {{incompatible payload operation name}}29 %pad = transform.get_producer_of_operand %matmul_padded[1]30 : (!transform.any_op) -> !transform.op<"tensor.pad">31 32 // We do not even reach this transform op.33 transform.structured.hoist_pad %pad by 1 loops34 : (!transform.op<"tensor.pad">) -> !transform.any_op35 transform.yield36 }37}38 39// -----40 41func.func @pad_and_hoist_init(42 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)43 -> tensor<24x25xf32>44{45 // expected-note @below {{when applied to this op}}46 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>47 func.return %0 : tensor<24x25xf32>48}49 50module attributes {transform.with_named_sequence} {51 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {52 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg153 : (!transform.any_op) -> !transform.any_op54 55 56 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)57 58 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {59 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],60 padding_dimensions=[0, 1, 2],61 copy_back_op = "none"62 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)63 64 %pad = transform.get_producer_of_operand %matmul_padded[2]65 : (!transform.any_op) -> !transform.op<"tensor.pad">66 67 // We do not know yet how to hoist the init.68 // expected-error @below {{transform.structured.hoist_pad failed to apply}}69 transform.structured.hoist_pad %pad by 1 loops70 : (!transform.op<"tensor.pad">) -> !transform.any_op71 transform.yield72 }73}74 75// -----76 77// CHECK-LABEL: pad_and_hoist_lhs(78func.func @pad_and_hoist_lhs(79 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)80 -> tensor<24x25xf32>81{82 // CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x5x12xf32>) {83 // CHECK: tensor.pad %{{.*}}84 // CHECK: : tensor<?x12xf32> to tensor<5x12xf32>85 // CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1]86 // CHECK-SAME: : tensor<5x12xf32> into tensor<5x5x12xf32>87 // CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {88 // CHECK: %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1]89 // CHECK-SAME: : tensor<5x5x12xf32> to tensor<5x12xf32>90 // CHECK: linalg.matmul ins(%[[PADDED]]91 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>92 func.return %0 : tensor<24x25xf32>93}94 95module attributes {transform.with_named_sequence} {96 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {97 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg198 : (!transform.any_op) -> !transform.any_op99 100 101 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)102 103 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {104 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],105 padding_dimensions=[0, 1, 2],106 copy_back_op = "none"107 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)108 109 %pad = transform.get_producer_of_operand %matmul_padded[0]110 : (!transform.any_op) -> !transform.any_op111 112 transform.structured.hoist_pad %pad by 1 loops113 : (!transform.any_op) -> !transform.any_op114 transform.yield115 }116}117 118// -----119 120// CHECK-LABEL: pad_and_hoist_lhs_transpose121func.func @pad_and_hoist_lhs_transpose(122 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)123 -> tensor<24x25xf32>124{125 // CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x12x5xf32>) {126 // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}}127 // CHECK: : tensor<?x12xf32> to tensor<5x12xf32>128 // CHECK: linalg.transpose129 // CHECK: ins(%[[PAD]] : tensor<5x12xf32>) outs(%{{.*}} : tensor<12x5xf32>)130 // CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1]131 // CHECK-SAME: : tensor<12x5xf32> into tensor<5x12x5xf32>132 // CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {133 // CHECK: %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1]134 // CHECK-SAME: : tensor<5x12x5xf32> to tensor<12x5xf32>135 // CHECK: %[[TRANSPOSED:.*]] = linalg.transpose ins(%[[PADDED]] : tensor<12x5xf32>)136 // CHECK: outs(%{{.*}} : tensor<5x12xf32>137 // CHECK: linalg.matmul ins(%[[TRANSPOSED]]138 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>139 func.return %0 : tensor<24x25xf32>140}141 142module attributes {transform.with_named_sequence} {143 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {144 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1145 : (!transform.any_op) -> !transform.any_op146 147 148 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)149 150 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {151 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],152 padding_dimensions=[0, 1, 2],153 copy_back_op = "none"154 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)155 156 %pad = transform.get_producer_of_operand %matmul_padded[0]157 : (!transform.any_op) -> !transform.any_op158 159 transform.structured.hoist_pad %pad by 1 loops, transpose by [1, 0]160 : (!transform.any_op) -> !transform.any_op161 transform.yield162 }163}164 165// -----166 167// CHECK-LABEL: pad_and_hoist_init168func.func @pad_and_hoist_init(169 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)170 -> tensor<24x25xf32>171{172 173 // CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {174 // CHECK: %[[PADDED:.*]] = tensor.pad %{{.*}}175 // CHECK: : tensor<?x25xf32> to tensor<5x25xf32>176 // CHECK: %[[SCF_YIELD:.*]] = scf.for %{{.*}} iter_args(%[[INNER_PADDED:[0-9a-zA-Z]*]] = %[[PADDED]]) -> (tensor<5x25xf32>)177 // CHECK: %[[RES:.*]] = linalg.matmul {{.*}} outs(%[[INNER_PADDED]]178 // CHECK-SAME: : tensor<5x25xf32>179 // CHECK: scf.yield %[[RES]] : tensor<5x25xf32>180 // CHECK: %[[EXTRACTED:.*]] = tensor.extract_slice %[[SCF_YIELD]][%{{.*}}, 0] [%{{.*}}, 25] [1, 1]181 // CHECK-SAME: : tensor<5x25xf32> to tensor<?x25xf32>182 // CHECK: tensor.insert_slice %[[EXTRACTED]] into %{{.*}}[%{{.*}}, 0] [%{{.*}}, 25] [1, 1]183 // CHECK-SAME: : tensor<?x25xf32> into tensor<24x25xf32>184 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>185 func.return %0 : tensor<24x25xf32>186}187 188module attributes {transform.with_named_sequence} {189 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {190 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1191 : (!transform.any_op) -> !transform.any_op192 193 194 %matmul_l1, %loops_l1:2 = transform.structured.tile_using_for %matmul tile_sizes [5, 0, 7] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)195 196 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {197 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],198 padding_dimensions=[0, 1, 2],199 copy_back_op = "none"200 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)201 202 %pad = transform.get_producer_of_operand %matmul_padded[2]203 : (!transform.any_op) -> !transform.op<"tensor.pad">204 205 transform.apply_licm to %loops_l1#1 : !transform.any_op206 207 transform.structured.hoist_pad %pad by 1 loops208 : (!transform.op<"tensor.pad">) -> !transform.any_op209 transform.yield210 }211}212