187 lines · plain
1// BUILD-PACKING-LOOP-NEST only checks the creation of packing code but does not connect it.2// Do not run canonicalization as it would be DCE'd away.3// RUN: mlir-opt --transform-interpreter -split-input-file --verify-diagnostics %s | FileCheck %s --check-prefix=BUILD-PACKING-LOOP-NEST4 5func.func @pad_and_hoist_rhs(6 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)7 -> tensor<24x25xf32>8{9 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>10 func.return %0 : tensor<24x25xf32>11}12 13module attributes {transform.with_named_sequence} {14 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {15 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg116 : (!transform.any_op) -> !transform.any_op17 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 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)24 25 // In this case, the pad op is actually empty: we only tile the first dimension26 // and it does not have an impact on the RHS operand.27 %pad = transform.get_producer_of_operand %matmul_padded[1]28 : (!transform.any_op) -> !transform.any_op29 30 // expected-error @below {{requires exactly 2 non-null handles}}31 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l132 : (!transform.any_op, !transform.any_op) -> !transform.any_op33 transform.yield34 }35}36 37// -----38 39func.func @pad_and_hoist_init(40 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)41 -> tensor<24x25xf32>42{43 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>44 func.return %0 : tensor<24x25xf32>45}46 47module attributes {transform.with_named_sequence} {48 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {49 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg150 : (!transform.any_op) -> !transform.any_op51 52 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)53 54 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {55 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],56 padding_dimensions=[0, 1, 2]57 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)58 59 %pad = transform.get_producer_of_operand %matmul_padded[2]60 : (!transform.any_op) -> !transform.any_op61 62 // We do not know yet how to hoist the init.63 // expected-error @below {{could not build packing loop nest}}64 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l165 : (!transform.any_op, !transform.any_op) -> !transform.any_op66 transform.yield67 }68}69 70// -----71 72// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_lhs73func.func @pad_and_hoist_lhs(74 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)75 -> tensor<24x25xf32>76{77 // BUILD-PACKING-LOOP-NEST: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<?x5x12xf32>) {78 // BUILD-PACKING-LOOP-NEST: tensor.pad %{{.*}} 79 // BUILD-PACKING-LOOP-NEST: : tensor<?x12xf32> to tensor<5x12xf32>80 // BUILD-PACKING-LOOP-NEST: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1]81 // BUILD-PACKING-LOOP-NEST-SAME: : tensor<5x12xf32> into tensor<?x5x12xf32>82 // BUILD-PACKING-LOOP-NEST: scf.for %{{.*}} -> (tensor<24x25xf32>)83 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>84 func.return %0 : tensor<24x25xf32>85}86 87module attributes {transform.with_named_sequence} {88 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {89 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg190 : (!transform.any_op) -> !transform.any_op91 92 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)93 94 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {95 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],96 padding_dimensions=[0, 1, 2]97 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)98 99 %pad = transform.get_producer_of_operand %matmul_padded[0]100 : (!transform.any_op) -> !transform.any_op101 102 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1103 : (!transform.any_op, !transform.any_op) -> !transform.any_op104 transform.yield105 }106}107 108// -----109 110// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_lhs_transpose111func.func @pad_and_hoist_lhs_transpose(112 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)113 -> tensor<24x25xf32>114{115 // BUILD-PACKING-LOOP-NEST: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<?x12x5xf32>) {116 // BUILD-PACKING-LOOP-NEST: tensor.pad %{{.*}}117 // BUILD-PACKING-LOOP-NEST: : tensor<?x12xf32> to tensor<5x12xf32>118 // BUILD-PACKING-LOOP-NEST: linalg.transpose119 // BUILD-PACKING-LOOP-NEST: ins({{.*}} : tensor<5x12xf32>) outs({{.*}} : tensor<12x5xf32>)120 // BUILD-PACKING-LOOP-NEST: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1]121 // BUILD-PACKING-LOOP-NEST-SAME: : tensor<12x5xf32> into tensor<?x12x5xf32>122 // BUILD-PACKING-LOOP-NEST: scf.for %{{.*}} -> (tensor<24x25xf32>)123 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>124 func.return %0 : tensor<24x25xf32>125}126 127module attributes {transform.with_named_sequence} {128 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {129 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1130 : (!transform.any_op) -> !transform.any_op131 132 %matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)133 134 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {135 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],136 padding_dimensions=[0, 1, 2]137 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)138 139 %pad = transform.get_producer_of_operand %matmul_padded[0]140 : (!transform.any_op) -> !transform.any_op141 142 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1, transpose by [1, 0]143 : (!transform.any_op, !transform.any_op) -> !transform.any_op144 transform.yield145 }146}147 148// -----149 150// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_init151func.func @pad_and_hoist_init(152 %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)153 -> tensor<24x25xf32>154{155 156 // BUILD-PACKING-LOOP-NEST: scf.for %{{.*}} -> (tensor<24x25xf32>) {157 // BUILD-PACKING-LOOP-NEST: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice158 // BUILD-PACKING-LOOP-NEST: %[[PADDED:.*]] = tensor.pad %[[EXTRACTED_SLICE]]159 // BUILD-PACKING-LOOP-NEST: : tensor<?x25xf32> to tensor<5x25xf32>160 // BUILD-PACKING-LOOP-NEST: scf.for %{{.*}} iter_args({{.*}} = %[[EXTRACTED_SLICE]]) -> (tensor<24x25xf32>, tensor<?x25xf32>) {161 %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>162 func.return %0 : tensor<24x25xf32>163}164 165module attributes {transform.with_named_sequence} {166 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {167 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1168 : (!transform.any_op) -> !transform.any_op169 170 %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)171 172 %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {173 padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],174 padding_dimensions=[0, 1, 2]175 } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)176 177 %pad = transform.get_producer_of_operand %matmul_padded[2]178 : (!transform.any_op) -> !transform.any_op179 180 transform.apply_licm to %loops_l1#1 : !transform.any_op181 182 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1#1183 : (!transform.any_op, !transform.any_op) -> !transform.any_op184 transform.yield185 }186}187