| // RUN: mlir-opt %s -split-input-file -linalg-fold-unit-extent-dims | FileCheck %s |
| |
| #accesses = [ |
| affine_map<(i, j, k, l, m) -> (i, k, m)>, |
| affine_map<(i, j, k, l, m) -> (i, k, j, l, m)> |
| ] |
| |
| #trait = { |
| iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"], |
| indexing_maps = #accesses, |
| library_call = "some_external_func" |
| } |
| |
| func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32> { |
| %0 = linalg.generic #trait |
| ins(%arg0 : tensor<?x1x?xf32>) |
| outs(%shape : tensor<?x1x?x1x?xf32>) { |
| ^bb0(%arg2 : f32, %arg3 : f32) : |
| linalg.yield %arg2 : f32 |
| } -> tensor<?x1x?x1x?xf32> |
| return %0 : tensor<?x1x?x1x?xf32> |
| } |
| // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> |
| // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d2)> |
| // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> |
| // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> |
| // CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)> |
| // CHECK-DAG: #[[$MAP5:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)> |
| // CHECK-DAG: #[[$MAP6:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d4)> |
| // CHECK-LABEL: func @drop_one_trip_loops |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]] |
| // CHECK: linalg.generic |
| // CHECK-SAME: indexing_maps = [#[[$MAP2]], #[[$MAP3]]] |
| // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"] |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[$MAP4]], #[[$MAP5]], #[[$MAP6]]] |
| |
| // ----- |
| |
| #accesses = [ |
| affine_map<(i, j, k, l, m) -> (i, k, m)>, |
| affine_map<(i, j, k, l, m) -> (i, k, j, l, m)> |
| ] |
| |
| #trait = { |
| iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"], |
| indexing_maps = #accesses, |
| library_call = "some_external_func" |
| } |
| |
| func @drop_one_trip_loops_indexed_generic |
| (%arg0 : tensor<?x1x?xi32>, %shape: tensor<?x1x?x1x?xi32>) -> tensor<?x1x?x1x?xi32> |
| { |
| %0 = linalg.indexed_generic #trait |
| ins(%arg0 : tensor<?x1x?xi32>) |
| outs(%shape: tensor<?x1x?x1x?xi32>) { |
| ^bb0(%arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, |
| %arg5 : index, %arg6 : i32, %arg7 : i32) : |
| %1 = addi %arg1, %arg2 : index |
| %2 = addi %1, %arg3 : index |
| %3 = addi %2, %arg4 : index |
| %4 = addi %3, %arg5 : index |
| %5 = index_cast %4 : index to i32 |
| %6 = addi %5, %arg6 : i32 |
| linalg.yield %6 : i32 |
| } -> tensor<?x1x?x1x?xi32> |
| return %0 : tensor<?x1x?x1x?xi32> |
| } |
| // CHECK-LABEL: func @drop_one_trip_loops_indexed_generic |
| // CHECK: linalg.indexed_generic |
| // CHECK: ^{{.+}}( |
| // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, %[[ARG2:[a-zA-Z0-9]+]]: index |
| // CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index, %[[ARG4:[a-zA-Z0-9]+]]: i32, %{{.*}}: i32) |
| // CHECK: %[[T3:.+]] = addi %[[ARG1]], %[[ARG2]] |
| // CHECK: %[[T4:.+]] = addi %[[T3]], %[[ARG3]] |
| // CHECK: %[[T5:.+]] = index_cast %[[T4]] : index to i32 |
| // CHECK: %[[T6:.+]] = addi %[[T5]], %[[ARG4]] : i32 |
| // CHECK: linalg.yield %[[T6]] : i32 |
| |
| // ----- |
| |
| #map0 = affine_map<(i, j) -> (i, j)> |
| #access = [#map0, #map0] |
| #trait = { |
| iterator_types = ["parallel", "parallel"], |
| indexing_maps = #access, |
| library_call = "some_external_func" |
| } |
| |
| func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32> |
| { |
| %0 = linalg.generic #trait |
| ins(%arg0 : tensor<1x1xf32>) |
| outs(%arg0 : tensor<1x1xf32>) { |
| ^bb0(%arg1: f32, %arg2: f32) : |
| linalg.yield %arg1 : f32 |
| } -> tensor<1x1xf32> |
| return %0 : tensor<1x1xf32> |
| } |
| // CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> ()> |
| // CHECK-LABEL: func @drop_all_loops |
| // CHECK: linalg.tensor_reshape %{{.*}} [] |
| // CHECK: linalg.generic |
| // CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]] |
| // CHECK-SAME: iterator_types = [] |
| |
| // ----- |
| |
| #map0 = affine_map<(i, j) -> (i, j)> |
| #access = [#map0, #map0] |
| #trait = { |
| iterator_types = ["parallel", "parallel"], |
| indexing_maps = #access, |
| library_call = "some_external_func" |
| } |
| |
| func @drop_all_loops_indexed_generic |
| (%arg0 : tensor<1x1xi32>) -> tensor<1x1xi32>{ |
| %0 = linalg.indexed_generic #trait |
| ins(%arg0 : tensor<1x1xi32>) |
| outs(%arg0 : tensor<1x1xi32>) { |
| ^bb0(%arg1 : index, %arg2 : index, %arg3: i32, %arg4: i32) : |
| %1 = addi %arg1, %arg2 : index |
| %2 = index_cast %1 : index to i32 |
| %3 = addi %2, %arg3 : i32 |
| linalg.yield %3 : i32 |
| } -> tensor<1x1xi32> |
| return %0 : tensor<1x1xi32> |
| } |
| |
| // CHECK-LABEL: func @drop_all_loops_indexed_generic |
| // CHECK: linalg.indexed_generic |
| // CHECK: ^{{.+}}(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32) |
| // CHECK: linalg.yield %[[ARG1]] : i32 |
| |
| // ----- |
| |
| #accesses = [ |
| affine_map<(d0) -> (0, d0)>, |
| affine_map<(d0) -> (d0)> |
| ] |
| |
| #trait = { |
| indexing_maps = #accesses, |
| iterator_types = ["parallel"], |
| library_call = "some_external_fn" |
| } |
| |
| func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> { |
| %0 = linalg.generic #trait |
| ins(%arg0 : tensor<1x5xf32>) |
| outs(%shape : tensor<5xf32>) { |
| ^bb0(%arg2: f32, %arg3: f32): // no predecessors |
| linalg.yield %arg2 : f32 |
| } -> tensor<5xf32> |
| return %0 : tensor<5xf32> |
| } |
| // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK-LABEL: func @leading_dim_1_canonicalization |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[$MAP0]]] |
| // CHECK: linalg.generic |
| // CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP1]]] |
| // CHECK-SAME: iterator_types = ["parallel"] |
| |
| // ----- |
| |
| #accesses = [ |
| affine_map<(d0, d1) -> (0, d1)>, |
| affine_map<(d0, d1) -> (d0, 0)>, |
| affine_map<(d0, d1) -> (d0, d1)> |
| ] |
| |
| #trait = { |
| indexing_maps = #accesses, |
| iterator_types = ["parallel", "parallel"], |
| library_call = "some_external_fn" |
| } |
| |
| func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>, %shape : tensor<5x5xf32>) -> tensor<5x5xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] : |
| tensor<5xf32> into tensor<1x5xf32> |
| %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] : |
| tensor<5xf32> into tensor<5x1xf32> |
| %2 = linalg.generic #trait |
| ins(%0, %1 : tensor<1x5xf32>, tensor<5x1xf32>) |
| outs(%shape : tensor<5x5xf32>) { |
| ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): |
| %3 = addf %arg3, %arg4 : f32 |
| linalg.yield %3 : f32 |
| } -> tensor<5x5xf32> |
| return %2 : tensor<5x5xf32> |
| } |
| // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)> |
| // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)> |
| // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK-LABEL: func @broadcast_test |
| // CHECK-NOT: linalg.tensor_reshape |
| // CHECK: linalg.generic |
| // CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]] |
| // CHECK-SAME: iterator_types = ["parallel", "parallel"] |
| // CHECK-NOT: linalg.tensor_reshape |
| |
| // ----- |
| |
| #accesses = [ |
| affine_map<(d0, d1) -> (0, 0)>, |
| affine_map<(d0, d1) -> (d0, d1)> |
| ] |
| |
| #trait = { |
| indexing_maps = #accesses, |
| iterator_types = ["parallel", "parallel"], |
| library_call = "some_external_fn" |
| } |
| |
| func @broadcast_scalar(%arg0 : tensor<1x1xf32>, %shape : tensor<?x?xf32>) -> tensor<?x?xf32> |
| { |
| %0 = linalg.generic #trait |
| ins(%arg0 : tensor<1x1xf32>) |
| outs(%shape : tensor<?x?xf32>) { |
| ^bb0(%arg2 : f32, %arg3 : f32): |
| linalg.yield %arg2 : f32 |
| } -> tensor<?x?xf32> |
| return %0 : tensor<?x?xf32> |
| } |
| // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()> |
| // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK-LABEL: func @broadcast_scalar |
| // CHECK-SAME: %[[ARG0:.*]]: tensor<1x1xf32> |
| // CHECK: %[[A:.*]] = linalg.tensor_reshape %[[ARG0]] [] |
| // CHECK-SAME: tensor<1x1xf32> into tensor<f32> |
| // CHECK: linalg.generic |
| // CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]] |
| // CHECK-SAME: iterator_types = ["parallel", "parallel"] |
| // CHECK-SAME: %[[A]] |
| |
| // ----- |
| |
| // CHECK: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK: func @fold_reshape |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]]] |
| // CHECK-SAME: tensor<2048xf32> into tensor<4x512xf32> |
| func @fold_reshape(%arg0 : tensor<2048xf32>) -> tensor<4x512xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 |
| [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>] |
| : tensor<2048xf32> into tensor<1x4x1x512xf32> |
| %1 = linalg.tensor_reshape %0 |
| [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>, |
| affine_map<(d0, d1, d2, d3) -> (d3)>] |
| : tensor<1x4x1x512xf32> into tensor<4x512xf32> |
| return %1 : tensor<4x512xf32> |
| } |
| |
| // ----- |
| |
| // CHECK: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK: func @fold_reshape |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]]] |
| // CHECK-SAME: tensor<4x512xf32> into tensor<2048xf32> |
| func @fold_reshape(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 |
| [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>, |
| affine_map<(d0, d1, d2, d3) -> (d3)>] |
| : tensor<4x512xf32> into tensor<1x4x1x512xf32> |
| %1 = linalg.tensor_reshape %0 |
| [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>] |
| : tensor<1x4x1x512xf32> into tensor<2048xf32> |
| return %1 : tensor<2048xf32> |
| } |
| |
| // ----- |
| |
| // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> |
| // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)> |
| // CHECK: func @fold_reshape |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]], #[[MAP1]]] |
| // CHECK-SAME: tensor<2048x1xf32> into tensor<4x512x1xf32> |
| func @fold_reshape(%arg0 : tensor<2048x1xf32>) -> tensor<4x512x1xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 |
| [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>, |
| affine_map<(d0, d1, d2, d3, d4) -> (d4)>] |
| : tensor<2048x1xf32> into tensor<1x4x1x512x1xf32> |
| %1 = linalg.tensor_reshape %0 |
| [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>, |
| affine_map<(d0, d1, d2, d3, d4) -> (d3)>, |
| affine_map<(d0, d1, d2, d3, d4) -> (d4)>] |
| : tensor<1x4x1x512x1xf32> into tensor<4x512x1xf32> |
| return %1 : tensor<4x512x1xf32> |
| } |
| |
| // ----- |
| |
| // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)> |
| // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2)> |
| // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)> |
| // CHECK: func @fold_reshape |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]], #[[MAP1]], #[[MAP2]]] |
| // CHECK-SAME: tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32> |
| func @fold_reshape(%arg0 : tensor<2048x1x2048xf32>) -> tensor<4x512x1x512x4xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 |
| [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d4)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d5)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d6, d7, d8)>] |
| : tensor<2048x1x2048xf32> into tensor<1x4x1x512x1x1x512x1x4xf32> |
| %1 = linalg.tensor_reshape %0 |
| [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d3, d4)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d5)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d6, d7)>, |
| affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d8)>] |
| : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32> |
| return %1 : tensor<4x512x1x512x4xf32> |
| } |
| |
| // ----- |
| |
| // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> |
| // CHECK: func @fold_reshape |
| // CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]] |
| // CHECK-SAME: tensor<2xf32> into tensor<2x1xf32> |
| func @fold_reshape(%arg0: tensor<2xf32>) -> tensor<2x1xf32> |
| { |
| %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2) -> (d0, d1, d2)>] : tensor<2xf32> into tensor<2x1x1xf32> |
| %1 = linalg.tensor_reshape %0 |
| [affine_map<(d0, d1, d2) -> (d0)>, |
| affine_map<(d0, d1, d2) -> (d1, d2)> |
| ] : tensor<2x1x1xf32> into tensor<2x1xf32> |
| return %1 : tensor<2x1xf32> |
| } |
| |
| // ----- |
| |
| #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> |
| #map1 = affine_map<(d0, d1, d2) -> (d2)> |
| #map3 = affine_map<(d0, d1, d2) -> (d0, d1)> |
| #map4 = affine_map<(d0, d1, d2) -> (d2)> |
| func @fold_unit_dim_tensor_reshape_op(%arg0 : tensor<5xf32>) -> tensor<2x5xf32> |
| { |
| %1 = linalg.init_tensor [1, 2, 5] : tensor<1x2x5xf32> |
| %2 = linalg.generic {i64, indexing_maps = [#map1, #map0], |
| iterator_types = ["parallel", "parallel", "parallel"]} |
| ins(%arg0 : tensor<5xf32>) outs(%1 : tensor<1x2x5xf32>) { |
| ^bb0(%arg1: f32, %arg2: f32): // no predecessors |
| linalg.yield %arg1 : f32 |
| } -> tensor<1x2x5xf32> |
| %3 = linalg.tensor_reshape %2 [#map3, #map4] |
| : tensor<1x2x5xf32> into tensor<2x5xf32> |
| return %3 : tensor<2x5xf32> |
| } |
| // CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op |
| // CHECK: %[[RESULT:.+]] = linalg.generic |
| // CHECK: return %[[RESULT]] |
