1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s 2 3// CHECK-LABEL: func @matmul_tensors( 4// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> 5// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32> 6// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { 7func.func @matmul_tensors( 8 %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) 9 -> tensor<?x?xf32> { 10// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) { 11// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) { 12// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) { 13// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 14// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 15// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 16// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>) 17// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32> 18// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32> 19// CHECK: scf.yield %[[TD]] : tensor<?x?xf32> 20// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32> 21// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32> 22 %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) 23 outs(%arg2: tensor<?x?xf32>) 24 -> tensor<?x?xf32> 25 26// CHECK: return %[[TD0]] : tensor<?x?xf32> 27 return %0 : tensor<?x?xf32> 28} 29 30module attributes {transform.with_named_sequence} { 31 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 32 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op 33 %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [2, 3, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 34 transform.yield 35 } 36} 37 38// ----- 39 40// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> 41// CHECK-NEXT: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)> 42// CHECK-NEXT: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> 43#access_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, 44 affine_map<(d0, d1, d2) -> (d2, d1)>, 45 affine_map<(d0, d1, d2) -> (d0, d1)>] 46 47// CHECK-LABEL: func @matmul_as_contract_tensors( 48// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> 49// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32> 50// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { 51func.func @matmul_as_contract_tensors( 52 %A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) 53 -> tensor<?x?xf32> { 54// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) { 55// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) { 56// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) { 57// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 58// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 59// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 60// CHECK: %[[sTD:.*]] = linalg.contract 61// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]] 62// CHECK-SAME: ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>) 63// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32> 64// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32> 65// CHECK: scf.yield %[[TD]] : tensor<?x?xf32> 66// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32> 67// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32> 68 %0 = linalg.contract indexing_maps = #access_maps 69 ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>) 70 outs(%C: tensor<?x?xf32>) 71 -> tensor<?x?xf32> 72 73// CHECK: return %[[TD0]] : tensor<?x?xf32> 74 return %0 : tensor<?x?xf32> 75} 76 77module attributes {transform.with_named_sequence} { 78 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 79 %0 = transform.structured.match ops{["linalg.contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op 80 %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [2, 3, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 81 transform.yield 82 } 83} 84 85// ----- 86 87// CHECK-LABEL: func @matmul_tensors_with_size_zeros( 88// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> 89// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32> 90// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { 91func.func @matmul_tensors_with_size_zeros( 92 %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) 93 -> tensor<?x?xf32> { 94 95// CHECK: %[[RES:.*]] = linalg.matmul ins(%[[TA]], %[[TB]] : tensor<?x?xf32>, tensor<?x?xf32>) 96// CHECK-SAME: outs(%[[TC]] : tensor<?x?xf32>) -> tensor<?x?xf32> 97// CHECK: return %[[RES]] 98 %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) 99 outs(%arg2: tensor<?x?xf32>) 100 -> tensor<?x?xf32> 101 return %0 : tensor<?x?xf32> 102} 103 104module attributes {transform.with_named_sequence} { 105 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 106 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op 107 %1 = transform.structured.tile_using_for %0 tile_sizes [0, 0, 0] : (!transform.any_op) -> (!transform.any_op) 108 transform.yield 109 } 110} 111 112// ----- 113 114func.func @generic_op_tensors( 115 %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { 116 %c0 = arith.constant 0 : index 117 %c1 = arith.constant 1 : index 118 %c2 = arith.constant 2 : index 119 %0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> 120 %1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> 121 %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> 122 %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32> 123 %4 = linalg.generic 124 {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, 125 affine_map<(d0, d1, d2) -> (d0, d2, d1)>, 126 affine_map<(d0, d1, d2) -> (d2, d1, d0)>], 127 iterator_types = ["parallel", "parallel", "parallel"]} 128 ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) 129 outs(%3 : tensor<?x?x?xf32>) { 130 ^bb0(%arg2 : f32, %arg3: f32, %arg4: f32): 131 %5 = arith.addf %arg2, %arg3 : f32 132 linalg.yield %5 : f32 133 } -> tensor<?x?x?xf32> 134 return %4 : tensor<?x?x?xf32> 135} 136 137module attributes {transform.with_named_sequence} { 138 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 139 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op 140 %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [2, 3, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 141 transform.yield 142 } 143} 144 145// CHECK-LABEL: func @generic_op_tensors 146// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 147// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 148// CHECK: %[[INIT:.+]] = tensor.empty 149// CHECK: %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) { 150// CHECK: %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) { 151// CHECK: %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) { 152// CHECK: %[[STARG0:.+]] = tensor.extract_slice %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 153// CHECK: %[[STARG1:.+]] = tensor.extract_slice %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 154// CHECK: %[[STARG2:.+]] = tensor.extract_slice %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 155// CHECK: %[[STRETURN:.+]] = linalg.generic 156// CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) 157// CHECK-SAME: outs(%[[STARG2]] : tensor<?x?x?xf32>) 158// CHECK: %[[TD:.+]] = tensor.insert_slice %[[STRETURN]] into %[[TC2]] 159// CHECK: scf.yield %[[TD]] 160// CHECK: } 161// CHECK: scf.yield %[[TD2]] 162// CHECK: } 163// CHECK: scf.yield %[[TD1]] 164// CHECK: } 165// CHECK: return %[[TD0]] 166 167// ----- 168 169// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)> 170 171// CHECK: fold_extract_slice 172// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<?x128xf32> 173// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<?x42xf32> 174func.func @fold_extract_slice( 175 %arg0 : tensor<?x128xf32>, %arg1 : tensor<?x42xf32>, %arg2 : tensor<?x42x?xf32>) -> tensor<?x42xf32> { 176 177 // CHECK: %[[C0:.*]] = arith.constant 0 178 %c0 = arith.constant 0 : index 179 180 // CHECK: %[[DIM:.*]] = tensor.dim %[[ARG1]], %[[C0]] 181 %0 = tensor.dim %arg1, %c0 : tensor<?x42xf32> 182 %1 = tensor.extract_slice %arg0[3, 4] [%0, 42] [1, 1] : tensor<?x128xf32> to tensor<?x42xf32> 183 184 // CHECK: %[[E:.*]] = tensor.extract_slice %[[ARG0]][3, 4] [%[[DIM]], 42] [1, 1] : tensor<?x128xf32> to tensor<?x42xf32> 185 186 // CHECK: scf.for %[[IV0:[0-9a-zA-Z]*]] = 187 // CHECK: scf.for %[[IV1:[0-9a-zA-Z]*]] = 188 189 // CHECK: %[[SIZE0:.*]] = affine.min #[[MAP0]](%[[IV0]])[%[[DIM]] 190 // Fold the existing extract slice op into the one created by the tiling. 191 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[E]] 192 // CHECK-SAME: %[[IV0]], %[[IV1]] 193 // CHECK-SAME: %[[SIZE0]], 3 194 // CHECK-SAME: 1, 1 195 // CHECK: {{.*}} = linalg.generic {{.*}} ins(%[[T0]] 196 %2 = linalg.generic 197 {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, 198 affine_map<(d0, d1, d2) -> (d0, d1, d2)>, 199 affine_map<(d0, d1, d2) -> (d0, d1)>], 200 iterator_types = ["parallel", "parallel", "parallel"]} 201 ins(%1, %arg2 : tensor<?x42xf32>, tensor<?x42x?xf32>) 202 outs(%arg1 : tensor<?x42xf32>) { 203 ^bb0(%arg3 : f32, %arg4: f32, %arg5: f32): 204 %5 = arith.addf %arg3, %arg5 : f32 205 linalg.yield %5 : f32 206 } -> tensor<?x42xf32> 207 return %2 : tensor<?x42xf32> 208} 209 210module attributes {transform.with_named_sequence} { 211 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 212 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op 213 %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [2, 3, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 214 transform.yield 215 } 216} 217