You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
329 lines
11 KiB
329 lines
11 KiB
// RUN: mlir-opt %s -linalg-fold-unit-extent-dims -split-input-file | 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>) -> tensor<?x1x?x1x?xf32>
|
|
{
|
|
%0 = linalg.generic #trait
|
|
ins(%arg0 : tensor<?x1x?xf32>) {
|
|
^bb0(%arg1 : f32) :
|
|
linalg.yield %arg1 : 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>) -> tensor<?x1x?x1x?xi32>
|
|
{
|
|
%0 = linalg.indexed_generic #trait
|
|
ins(%arg0 : tensor<?x1x?xi32>) {
|
|
^bb0(%arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index,
|
|
%arg5 : index, %arg6 : 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)
|
|
// 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>) {
|
|
^bb0(%arg1: 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>) {
|
|
^bb0(%arg1 : index, %arg2 : index, %arg3: 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)
|
|
// 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>) -> tensor<5xf32> {
|
|
%0 = linalg.generic #trait
|
|
ins(%arg0 : tensor<1x5xf32>) {
|
|
^bb0(%arg2: 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>) -> 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>) {
|
|
^bb0(%arg2: f32, %arg3: f32):
|
|
%3 = addf %arg2, %arg3 : 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>) -> tensor<?x?xf32>
|
|
{
|
|
%0 = linalg.generic #trait
|
|
ins(%arg0 : tensor<1x1xf32>) {
|
|
^bb0(%arg1 : f32):
|
|
linalg.yield %arg1 : 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>
|
|
}
|