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.
441 lines
21 KiB
441 lines
21 KiB
// RUN: mlir-opt %s -linalg-fusion-for-tensor-ops -split-input-file | FileCheck %s
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
|
|
func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?x?xf32>,
|
|
%arg1 : tensor<?x?x?xf32>) ->
|
|
tensor<?x?x?xf32>
|
|
{
|
|
%0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k)>,
|
|
affine_map<(i, j, k, l) -> (l)>] :
|
|
tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
|
|
%1 = linalg.generic {
|
|
indexing_maps = [#map0, #map1, #map1],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
|
|
^bb0(%arg3: f32, %arg4: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?x?xf32>
|
|
return %1 : tensor<?x?x?xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>
|
|
// CHECK: func @generic_op_reshape_producer_fusion
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?x?xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
|
|
// CHECK: %[[T1:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP3]], #[[MAP4]], #[[MAP4]]]
|
|
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[ARG0]], %[[T0]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
// CHECK: %[[T2:.+]] = linalg.tensor_reshape
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK-SAME: tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
|
|
// CHECK: return %[[T2]]
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1) -> (d0, d1)>
|
|
func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,
|
|
%arg1 : tensor<?x?xf32>) ->
|
|
tensor<?x?x4x5xf32>
|
|
{
|
|
%0 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
^bb0(%arg3: f32, %arg4: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?xf32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?xf32> into tensor<?x?x4x5xf32>
|
|
return %1 : tensor<?x?x4x5xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK: func @generic_op_reshape_consumer_fusion
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32>
|
|
// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32>
|
|
// CHECK: %[[T2:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]
|
|
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>)
|
|
// CHECK: return %[[T2]] : tensor<?x?x4x5xf32>
|
|
|
|
|
|
// -----
|
|
|
|
func @reshape_as_consumer_permutation
|
|
(%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>)
|
|
-> tensor<?x?x?x?x?x?xf32> {
|
|
%c = linalg.generic {
|
|
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
|
|
affine_map<(d0, d1, d2) -> (d1, d2)>,
|
|
affine_map<(d0, d1, d2) -> (d0, d2, d1)>],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>) {
|
|
^bb0(%arg0 : f32, %arg1: f32):
|
|
%1 = addf %arg0, %arg1 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%d = linalg.tensor_reshape %c
|
|
[affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
|
|
affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
|
|
affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>]
|
|
: tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
|
|
return %d : tensor<?x?x?x?x?x?xf32>
|
|
}
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
|
|
// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
|
|
// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
|
|
// CHECK: func @reshape_as_consumer_permutation
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>
|
|
// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
|
|
// CHECK-SAME: [#[[MAP3]], #[[MAP4]]]
|
|
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x?x?xf32>
|
|
// CHECK: %[[T2:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]]]
|
|
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x?x?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
// CHECK: return %[[T2]] : tensor<?x?x?x?x?x?xf32>
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1) -> (d0, d1)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d2)>
|
|
|
|
func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>)
|
|
-> tensor<8x33x4xf32> {
|
|
%cst = constant dense<2.000000e+00> : tensor<264x4xf32>
|
|
%0 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>) {
|
|
^bb0(%arg1: f32, %arg2: f32): // no predecessors
|
|
%2 = mulf %arg1, %arg2 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<264x4xf32>
|
|
%1 = linalg.tensor_reshape %0 [#map1, #map2] :
|
|
tensor<264x4xf32> into tensor<8x33x4xf32>
|
|
return %1 : tensor<8x33x4xf32>
|
|
}
|
|
|
|
// 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, d1, d2)>
|
|
// CHECK: func @generic_op_reshape_consumer_static
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: tensor<264x4xf32> into tensor<8x33x4xf32>
|
|
// CHECK: %[[T1:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]]]
|
|
// CHECK-SAME: ["parallel", "parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[T0]] : tensor<8x33x4xf32>)
|
|
// CHECK: return %[[T1]] : tensor<8x33x4xf32>
|
|
|
|
// -----
|
|
|
|
func @scalar_reshape(%arg0 : tensor<1x10xf32>, %arg1 : tensor<1xf32>)
|
|
-> tensor<1x10xf32> {
|
|
%0 = linalg.tensor_reshape %arg1 [] : tensor<1xf32> into tensor<f32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
|
|
iterator_types = ["parallel"]} ins(%0 : tensor<f32>) {
|
|
^bb0(%arg2: f32): // no predecessors
|
|
linalg.yield %arg2 : f32
|
|
} -> tensor<10xf32>
|
|
%2 = linalg.tensor_reshape %1 [affine_map<(d0, d1) -> (d0, d1)>]
|
|
: tensor<10xf32> into tensor<1x10xf32>
|
|
return %2 : tensor<1x10xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
|
|
// CHECK: func @scalar_reshape
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x10xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<1xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]] []
|
|
// CHECK-SAME: tensor<1xf32> into tensor<f32>
|
|
// CHECK: %[[T1:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[T0]] : tensor<f32>)
|
|
// CHECK: return %[[T1]] : tensor<1x10xf32>
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
|
|
func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>,
|
|
%arg1 : tensor<?x?x?xi32>) ->
|
|
tensor<?x?x?xi32>
|
|
{
|
|
%0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k)>,
|
|
affine_map<(i, j, k, l) -> (l)>] :
|
|
tensor<?x?x4x?xi32> into tensor<?x?x?xi32>
|
|
%1 = linalg.indexed_generic {
|
|
indexing_maps = [#map0, #map1, #map1],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>) {
|
|
^bb0(%arg3 : index, %arg4 : index, %arg5 : index, %arg6: i32, %arg7: i32):
|
|
%1 = muli %arg6, %arg7 : i32
|
|
%2 = index_cast %arg3 : index to i32
|
|
%3 = addi %1, %2 : i32
|
|
%4 = index_cast %arg4 : index to i32
|
|
%5 = addi %3, %4 : i32
|
|
%6 = index_cast %arg5 : index to i32
|
|
%7 = addi %5, %6 : i32
|
|
linalg.yield %7 : i32
|
|
} -> tensor<?x?x?xi32>
|
|
return %1 : tensor<?x?x?xi32>
|
|
}
|
|
|
|
// The generic op version of the test check for the op structure. Only
|
|
// checking the op body here.
|
|
// CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
|
|
// CHECK: func @indexed_generic_op_reshape_producer_fusion
|
|
// CHECK: linalg.indexed_generic
|
|
// CHECK: ^{{.*}}(
|
|
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32)
|
|
// CHECK: %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG2]], %[[ARG3]])
|
|
// CHECK: %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
|
|
// CHECK: %[[T5:.+]] = index_cast %[[T3]]
|
|
// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]]
|
|
// CHECK: %[[T7:.+]] = index_cast %[[ARG4]]
|
|
// CHECK: %[[T8:.+]] = addi %[[T6]], %[[T7]]
|
|
// CHECK: %[[T9:.+]] = index_cast %[[ARG5]]
|
|
// CHECK: %[[T10:.+]] = addi %[[T8]], %[[T9]]
|
|
// CHECK: linalg.yield %[[T10]]
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1) -> (d0, d1)>
|
|
func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>,
|
|
%arg1 : tensor<?x?xi32>) ->
|
|
tensor<?x?x4x5xi32>
|
|
{
|
|
%0 = linalg.indexed_generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>) {
|
|
^bb0(%arg3 : index, %arg4 : index, %arg5: i32, %arg6: i32): // no predecessors
|
|
%1 = muli %arg5, %arg6 : i32
|
|
%2 = index_cast %arg3 : index to i32
|
|
%3 = addi %1, %2 : i32
|
|
%4 = index_cast %arg4 : index to i32
|
|
%5 = addi %3, %4 : i32
|
|
linalg.yield %5 : i32
|
|
} -> tensor<?x?xi32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?xi32> into tensor<?x?x4x5xi32>
|
|
return %1 : tensor<?x?x4x5xi32>
|
|
}
|
|
// The generic op version of the test check for the op structure. Only
|
|
// checking the op body here.
|
|
// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * 20 + d1 * 5 + d2)>
|
|
// CHECK: func @indexed_generic_op_reshape_consumer_fusion
|
|
// CHECK: linalg.indexed_generic
|
|
// CHECK: ^{{.*}}(
|
|
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32)
|
|
// CHECK: %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG3]], %[[ARG4]], %[[ARG5]])
|
|
// CHECK: %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
|
|
// CHECK: %[[T5:.+]] = index_cast %[[ARG2]]
|
|
// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]]
|
|
// CHECK: %[[T7:.+]] = index_cast %[[T3]]
|
|
// CHECK: %[[T8:.+]] = addi %[[T6]], %[[T7]]
|
|
// CHECK: linalg.yield %[[T8]]
|
|
|
|
// -----
|
|
|
|
func @reshape_as_consumer_permutation
|
|
(%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>)
|
|
-> tensor<2x3x4x5x6x7xi32> {
|
|
%c = linalg.indexed_generic {
|
|
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
|
|
affine_map<(d0, d1, d2) -> (d1, d2)>,
|
|
affine_map<(d0, d1, d2) -> (d0, d2, d1)>],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%a, %b : tensor<210x6x4xi32>, tensor<210x4xi32>) {
|
|
^bb0(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : i32, %arg4: i32):
|
|
%1 = addi %arg3, %arg4 : i32
|
|
%2 = index_cast %arg0 : index to i32
|
|
%3 = addi %1, %2 : i32
|
|
%4 = index_cast %arg1 : index to i32
|
|
%5 = addi %3, %4 : i32
|
|
%6 = index_cast %arg2 : index to i32
|
|
%7 = addi %5, %6 : i32
|
|
linalg.yield %7 : i32
|
|
} -> tensor<6x4x210xi32>
|
|
%d = linalg.tensor_reshape %c
|
|
[affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
|
|
affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
|
|
affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>]
|
|
: tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>
|
|
return %d : tensor<2x3x4x5x6x7xi32>
|
|
}
|
|
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1) -> (d0 * 3 + d1)>
|
|
// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2) -> (d0 * 42 + d1 * 7 + d2)>
|
|
// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
|
|
// CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
|
|
// CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
|
|
// CHECK: func @reshape_as_consumer_permutation
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<210x6x4xi32>
|
|
// CHECK-SAME: %[[ARG1:.+]]: tensor<210x4xi32>
|
|
// CHECK-DAG: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK-DAG: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
|
|
// CHECK-SAME: [#[[MAP3]], #[[MAP4]]]
|
|
// CHECK: %[[T2:.+]] = linalg.indexed_generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP7]], #[[MAP8]], #[[MAP9]]]
|
|
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<{{.+}}>, tensor<{{.+}}>)
|
|
// CHECK: ^{{.+}}(
|
|
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index, %[[ARG7:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG8:[a-zA-Z0-9]+]]: i32, %[[ARG9:[a-zA-Z0-9]+]]: i32)
|
|
// CHECK-DAG: %[[T3:.+]] = affine.apply #[[MAP5]](%[[ARG2]], %[[ARG3]])
|
|
// CHECK-DAG: %[[T4:.+]] = affine.apply #[[MAP6]](%[[ARG4]], %[[ARG5]], %[[ARG6]])
|
|
// CHECK-DAG: %[[T5:.+]] = addi %[[ARG8]], %[[ARG9]]
|
|
// CHECK: %[[T6:.+]] = index_cast %[[T3]]
|
|
// CHECK: %[[T7:.+]] = addi %[[T5]], %[[T6]]
|
|
// CHECK: %[[T8:.+]] = index_cast %[[T4]]
|
|
// CHECK: %[[T9:.+]] = addi %[[T7]], %[[T8]]
|
|
// CHECK: %[[T10:.+]] = index_cast %[[ARG7]]
|
|
// CHECK: %[[T11:.+]] = addi %[[T9]], %[[T10]]
|
|
|
|
// -----
|
|
|
|
func @reshape_as_producer_projected_permutation
|
|
(%arg0 : tensor<33x8x?xi32>) -> tensor<264x?x4xi32> {
|
|
%0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2) -> (d0, d1)>,
|
|
affine_map<(d0, d1, d2) -> (d2)>]
|
|
: tensor<33x8x?xi32> into tensor<264x?xi32>
|
|
%1 = linalg.indexed_generic
|
|
{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>,
|
|
affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
|
|
iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<264x?xi32>) {
|
|
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: i32): // no predecessors
|
|
%2 = index_cast %arg1 : index to i32
|
|
%3 = addi %arg4, %2 : i32
|
|
%4 = index_cast %arg2 : index to i32
|
|
%5 = addi %3, %4 : i32
|
|
%6 = index_cast %arg3 : index to i32
|
|
%7 = addi %5, %6 : i32
|
|
linalg.yield %7 : i32
|
|
} -> tensor<264x?x4xi32>
|
|
return %1 : tensor<264x?x4xi32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0 * 8 + d1)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
|
|
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)>
|
|
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK: @reshape_as_producer_projected_permutation
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<33x8x?xi32>
|
|
// CHECK: %[[RES:.+]] = linalg.indexed_generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: ins(%[[ARG0]] : tensor<33x8x?xi32>)
|
|
// CHECK: ^{{.+}}(
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index,
|
|
// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: i32)
|
|
// CHECK: %[[T0:.+]] = affine.apply #[[MAP2]](%[[ARG1]], %[[ARG2]])
|
|
// CHECK: %[[T1:.+]] = index_cast %[[T0]] : index to i32
|
|
// CHECK: %[[T2:.+]] = addi %[[ARG5]], %[[T1]] : i32
|
|
// CHECK: %[[T3:.+]] = index_cast %[[ARG3]] : index to i32
|
|
// CHECK: %[[T4:.+]] = addi %[[T2]], %[[T3]] : i32
|
|
// CHECK: %[[T5:.+]] = index_cast %[[ARG4]] : index to i32
|
|
// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]] : i32
|
|
// CHECK: linalg.yield %[[T6]] : i32
|
|
// CHECK: %[[RES2:.+]] = linalg.tensor_reshape %[[RES]]
|
|
// CHECK-SAME: [#[[MAP3]], #[[MAP4]], #[[MAP5]]]
|
|
// CHECK-SAME: : tensor<33x8x?x4xi32> into tensor<264x?x4xi32>
|
|
// CHECK: return %[[RES2]] : tensor<264x?x4xi32>
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1) -> (d0, d1)>
|
|
#map1 = affine_map<(d0, d1) -> (d1, d0)>
|
|
func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>,
|
|
%arg1 : tensor<?x?xf32>) ->
|
|
tensor<?x?x4x5xf32>
|
|
{
|
|
%0 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map1],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
^bb0(%arg3: f32, %arg4: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?xf32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?xf32> into tensor<?x?x4x5xf32>
|
|
return %1 : tensor<?x?x4x5xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
|
|
// CHECK: func @generic_op_reshape_consumer_fusion_projected
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32>
|
|
// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32>
|
|
// CHECK: %[[T2:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]]]
|
|
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
|
|
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>)
|
|
// CHECK: return %[[T2]] : tensor<?x?x4x5xf32>
|