1// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" -split-input-file | FileCheck %s 2// RUN: mlir-opt %s -linalg-tile-to-tiled-loop="linalg-tile-sizes=2,3,4 linalg-distribution-types=block_x,block_y,none" -split-input-file | FileCheck %s -check-prefix=TLOOP 3 4// CHECK-LABEL: func @matmul_tensors( 5// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> 6// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32> 7// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { 8func @matmul_tensors( 9 %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) 10 -> tensor<?x?xf32> { 11// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) { 12// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) { 13// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) { 14// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 15// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 16// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> 17// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>) 18// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32> 19// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32> 20// CHECK: scf.yield %[[TD]] : tensor<?x?xf32> 21// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32> 22// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32> 23 %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) 24 outs(%arg2: tensor<?x?xf32>) 25 -> tensor<?x?xf32> 26 27// CHECK: return %[[TD0]] : tensor<?x?xf32> 28 return %0 : tensor<?x?xf32> 29} 30 31// TLOOP-LABEL: func @matmul_tensors 32// TLOOP-SAME: (%[[ARG_0:.*]]: [[TY:.*]], %[[ARG_1:.*]]: [[TY]], 33// TLOOP-SAME: %[[ARG_2:.*]]: [[TY]]) -> [[TY]] { 34 35// TLOOP-DAG: %[[C0:.*]] = constant 0 : index 36// TLOOP-DAG: %[[C1:.*]] = constant 1 : index 37// TLOOP-DAG: %[[C2:.*]] = constant 2 : index 38// TLOOP-DAG: %[[C3:.*]] = constant 3 : index 39// TLOOP-DAG: %[[C4:.*]] = constant 4 : index 40 41// TLOOP: %[[ARG_0_X:.*]] = tensor.dim %[[ARG_0]], %[[C0]] : [[TY]] 42// TLOOP: %[[ARG_0_Y:.*]] = tensor.dim %[[ARG_0]], %[[C1]] : [[TY]] 43// TLOOP: %[[ARG_1_Y:.*]] = tensor.dim %[[ARG_1]], %[[C1]] : [[TY]] 44 45// TLOOP: %{{.*}} = linalg.tiled_loop (%[[I:.*]], %[[J:.*]], %[[K:.*]]) = 46// TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]]) 47// TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_1_Y]], %[[ARG_0_Y]]) 48// TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]]) 49// TLOOP-SAME: ins (%[[A0:.*]] = %[[ARG_0]]: [[TY]], %[[A1:.*]] = %[[ARG_1]]: [[TY]]) 50// TLOOP-SAME: outs (%[[A2:.*]] = %[[ARG_2]]: [[TY]]) 51// TLOOP-SAME: iterators["parallel", "parallel", "reduction"] 52// TLOOP-SAME: distribution["block_x", "block_y", "none"] { 53 54// TLOOP: %[[SUB_ARG_0:.*]] = tensor.extract_slice %[[A0]][%[[I]], %[[K]]] 55// TLOOP: %[[SUB_ARG_1:.*]] = tensor.extract_slice %[[A1]][%[[K]], %[[J]]] 56// TLOOP: %[[SUB_ARG_2:.*]] = tensor.extract_slice %[[A2]][%[[I]], %[[J]]] 57 58// TLOOP: %[[PROD:.*]] = linalg.matmul ins(%[[SUB_ARG_0]], %[[SUB_ARG_1]] 59// TLOOP-SE: outs(%[[SUB_ARG_2]] : [[TY]]) -> [[TY]] 60 61// TLOOP: %[[O:.*]] = tensor.insert_slice %[[PROD]] into %[[A2]][%[[I]], %[[J]]] 62// TLOOP: linalg.yield %[[O]] : [[TY]] 63 64// ----- 65 66func @generic_op_tensors( 67 %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { 68 %c0 = constant 0 : index 69 %c1 = constant 1 : index 70 %c2 = constant 2 : index 71 %0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> 72 %1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> 73 %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> 74 %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> 75 %4 = linalg.generic 76 {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, 77 affine_map<(d0, d1, d2) -> (d0, d2, d1)>, 78 affine_map<(d0, d1, d2) -> (d2, d1, d0)>], 79 iterator_types = ["parallel", "parallel", "parallel"]} 80 ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) 81 outs(%3 : tensor<?x?x?xf32>) { 82 ^bb0(%arg2 : f32, %arg3: f32, %arg4: f32): 83 %5 = addf %arg2, %arg3 : f32 84 linalg.yield %5 : f32 85 } -> tensor<?x?x?xf32> 86 return %4 : tensor<?x?x?xf32> 87} 88 89// CHECK-LABEL: func @generic_op_tensors 90// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 91// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 92// CHECK: %[[INIT:.+]] = linalg.init_tensor 93// CHECK: %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) { 94// CHECK: %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) { 95// CHECK: %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) { 96// CHECK: %[[STARG0:.+]] = tensor.extract_slice %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 97// CHECK: %[[STARG1:.+]] = tensor.extract_slice %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 98// CHECK: %[[STARG2:.+]] = tensor.extract_slice %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> 99// CHECK: %[[STRETURN:.+]] = linalg.generic 100// CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) 101// CHECK-SAME: outs(%[[STARG2]] : tensor<?x?x?xf32>) 102// CHECK: %[[TD:.+]] = tensor.insert_slice %[[STRETURN]] into %[[TC2]] 103// CHECK: scf.yield %[[TD]] 104// CHECK: } 105// CHECK: scf.yield %[[TD2]] 106// CHECK: } 107// CHECK: scf.yield %[[TD1]] 108// CHECK: } 109// CHECK: return %[[TD0]] 110 111// TLOOP-LABEL: func @generic_op_tensors( 112// TLOOP-SAME: %[[ARG_0:.*]]: [[TY:.*]], 113// TLOOP-SAME: %[[ARG_1:.*]]: [[TY]]) -> [[TY]] { 114 115// TLOOP-DAG: %[[C0:.*]] = constant 0 : index 116// TLOOP-DAG: %[[C1:.*]] = constant 1 : index 117// TLOOP-DAG: %[[C2:.*]] = constant 2 : index 118// TLOOP-DAG: %[[C3:.*]] = constant 3 : index 119// TLOOP-DAG: %[[C4:.*]] = constant 4 : index 120 121// TLOOP: %[[INIT:.*]] = linalg.init_tensor 122// TLOOP: %[[ARG_0_X:.*]] = tensor.dim %[[ARG_0]], %[[C0]] : [[TY]] 123// TLOOP: %[[ARG_0_Y:.*]] = tensor.dim %[[ARG_0]], %[[C1]] : [[TY]] 124// TLOOP: %[[ARG_0_Z:.*]] = tensor.dim %[[ARG_0]], %[[C2]] : [[TY]] 125 126// TLOOP: %{{.*}} = linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = 127// TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]]) 128// TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_0_Y]], %[[ARG_0_Z]]) 129// TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]]) 130// TLOOP-SAME: ins (%{{.*}} = %[[ARG_0]]: [[TY]], %{{.*}} = %[[ARG_1]]: [[TY]]) 131// TLOOP-SAME: outs (%{{.*}} = %[[INIT]]: [[TY]]) 132// TLOOP-SAME: distribution["block_x", "block_y", "none"] { 133