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-NEST 4 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 %arg1 16 : (!transform.any_op) -> !transform.any_op 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 } : (!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 dimension 26 // 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_op 29 30 // expected-error @below {{requires exactly 2 non-null handles}} 31 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1 32 : (!transform.any_op, !transform.any_op) -> !transform.any_op 33 transform.yield 34 } 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 %arg1 50 : (!transform.any_op) -> !transform.any_op 51 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_op 61 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_l1 65 : (!transform.any_op, !transform.any_op) -> !transform.any_op 66 transform.yield 67 } 68} 69 70// ----- 71 72// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_lhs 73func.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 %arg1 90 : (!transform.any_op) -> !transform.any_op 91 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_op 101 102 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1 103 : (!transform.any_op, !transform.any_op) -> !transform.any_op 104 transform.yield 105 } 106} 107 108// ----- 109 110// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_lhs_transpose 111func.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.transpose 119 // 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 %arg1 130 : (!transform.any_op) -> !transform.any_op 131 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_op 141 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_op 144 transform.yield 145 } 146} 147 148// ----- 149 150// BUILD-PACKING-LOOP-NEST-LABEL: pad_and_hoist_init 151func.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_slice 158 // 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 %arg1 168 : (!transform.any_op) -> !transform.any_op 169 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_op 179 180 transform.apply_licm to %loops_l1#1 : !transform.any_op 181 182 transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1#1 183 : (!transform.any_op, !transform.any_op) -> !transform.any_op 184 transform.yield 185 } 186} 187