forked from OSchip/llvm-project
275 lines
12 KiB
MLIR
275 lines
12 KiB
MLIR
// RUN: mlir-opt -split-input-file -linalg-fold-reshape-ops-by-linearization -verify-diagnostics %s | FileCheck %s
|
|
|
|
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xf32>,
|
|
%arg1 : tensor<?x?x4x?xf32>) ->
|
|
tensor<?x?x4x?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?xf32> into tensor<?x?x4x?xf32>
|
|
%1 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
|
|
ins(%0, %arg1 : tensor<?x?x4x?xf32>, tensor<?x?x4x?xf32>)
|
|
outs(%0 : tensor<?x?x4x?xf32>) {
|
|
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?x4x?xf32>
|
|
return %1 : tensor<?x?x4x?xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
|
|
// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK: func @generic_op_reshape_producer_fusion
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP3]], #[[MAP4]], #[[MAP4]]]
|
|
// CHECK-SAME: ins(%[[ARG0]], %{{.+}} : tensor<?x?x?xf32>, tensor<?x?x4x?xf32>)
|
|
// CHECK-SAME: outs(%[[T0]] : tensor<?x?x4x?xf32>)
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xf32>,
|
|
%arg1 : tensor<?x?x4x5xf32>) ->
|
|
tensor<?x?xf32>
|
|
{
|
|
%0 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
|
|
ins(%arg0, %arg1 : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>)
|
|
outs(%arg0 : tensor<?x?x4x5xf32>){
|
|
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?x4x5xf32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?x4x5xf32> into tensor<?x?xf32>
|
|
return %1 : tensor<?x?xf32>
|
|
}
|
|
|
|
// 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-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
|
|
// CHECK: func @generic_op_reshape_consumer_fusion
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x5xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]]]
|
|
// CHECK-SAME: outs(%[[T0]] : tensor<?x?xf32>)
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xi32>)
|
|
-> tensor<?x?x4x?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?x?xi32> into tensor<?x?x4x?xi32>
|
|
%1 = linalg.indexed_generic {
|
|
indexing_maps = [#map0, #map0],
|
|
iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
|
|
ins(%0 : tensor<?x?x4x?xi32>)
|
|
outs(%0 : tensor<?x?x4x?xi32>) {
|
|
^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32, %arg7 : i32): // no predecessors
|
|
%2 = index_cast %arg2 : index to i32
|
|
%3 = addi %arg6, %2 : i32
|
|
linalg.yield %3 : i32
|
|
} -> tensor<?x?x4x?xi32>
|
|
return %1 : tensor<?x?x4x?xi32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
|
|
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK: func @indexed_generic_op_reshape_producer_fusion
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xi32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
|
|
// CHECK: linalg.indexed_generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP3]], #[[MAP4]]]
|
|
// CHECK-SAME: ins(%[[ARG0]] : tensor<?x?x?xi32>)
|
|
// CHECK-SAME: outs(%[[T0]] : tensor<?x?x4x?xi32>)
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
|
|
-> tensor<?x?xi32> {
|
|
%0 = linalg.indexed_generic {
|
|
indexing_maps = [#map0, #map0],
|
|
iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
|
|
ins(%arg0 : tensor<?x?x4x5xi32>) outs(%arg0 : tensor<?x?x4x5xi32>) {
|
|
^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32, %arg7: i32): // no predecessors
|
|
%2 = index_cast %arg2 : index to i32
|
|
%3 = addi %arg6, %2 : i32
|
|
linalg.yield %3 : i32
|
|
} -> tensor<?x?x4x5xi32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?x4x5xi32> into tensor<?x?xi32>
|
|
return %1 : tensor<?x?xi32>
|
|
}
|
|
|
|
// 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-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
|
|
// CHECK: func @indexed_generic_op_reshape_consumer_fusion
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x5xi32>
|
|
// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK: linalg.indexed_generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
|
|
// CHECK-SAME: outs(%[[T0]] : tensor<?x?xi32>)
|
|
// CHECK-NOT: linalg.tensor_reshape
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d0)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
|
|
#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
func @generic_op_021_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<3x7x5xf32> {
|
|
%0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
|
|
%1 = linalg.init_tensor [3, 7, 5] : tensor<3x7x5xf32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map2, #map3],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<3x7x5xf32>) {
|
|
^bb0(%arg2: f32, %arg3 : f32): // no predecessors
|
|
linalg.yield %arg2 : f32
|
|
} -> tensor<3x7x5xf32>
|
|
return %2 : tensor<3x7x5xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 + d2 * 7)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
// CHECK: func @generic_op_021_permultation_reshape_producer_fusion
|
|
// CHECK-NOT: linalg.tensor_reshape
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d0)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
|
|
#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
func @generic_op_120_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x7x3xf32> {
|
|
%0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
|
|
%1 = linalg.init_tensor [5, 7, 3] : tensor<5x7x3xf32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map2, #map3],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<5x7x3xf32>) {
|
|
^bb0(%arg2: f32, %arg3: f32): // no predecessors
|
|
linalg.yield %arg2 : f32
|
|
} -> tensor<5x7x3xf32>
|
|
return %2 : tensor<5x7x3xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0 * 7 + d1)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
// CHECK: func @generic_op_120_permultation_reshape_producer_fusion
|
|
// CHECK-NOT: linalg.tensor_reshape
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d0)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
|
|
#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
func @generic_op_102_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x3x7xf32> {
|
|
%0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
|
|
%1 = linalg.init_tensor [5, 3, 7] : tensor<5x3x7xf32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map2, #map3],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<5x3x7xf32>) {
|
|
^bb0(%arg2: f32, %arg3: f32): // no predecessors
|
|
linalg.yield %arg2 : f32
|
|
} -> tensor<5x3x7xf32>
|
|
return %2 : tensor<5x3x7xf32>
|
|
}
|
|
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
// CHECK: func @generic_op_102_permultation_reshape_producer_fusion
|
|
// CHECK-NOT: linalg.tensor_reshape
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0)>
|
|
#map3 = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
func @generic_op_102_permultation_reshape_consumer_fusion(%arg0 : tensor<3x5x7xf32>) -> tensor<5x21xf32> {
|
|
%0 = linalg.init_tensor [5, 3, 7] : tensor<5x3x7xf32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map0, #map1],
|
|
iterator_types = ["parallel", "parallel", "parallel"]}
|
|
ins(%arg0 : tensor<3x5x7xf32>) outs(%0 : tensor<5x3x7xf32>) {
|
|
^bb0(%arg2: f32, %arg3 : f32): // no predecessors
|
|
linalg.yield %arg2 : f32
|
|
} -> tensor<5x3x7xf32>
|
|
%2 = linalg.tensor_reshape %1 [#map2, #map3] : tensor<5x3x7xf32> into tensor<5x21xf32>
|
|
return %2 : tensor<5x21xf32>
|
|
}
|
|
|
|
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0)>
|
|
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
|
|
// CHECK: func @generic_op_102_permultation_reshape_consumer_fusion
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<3x5x7xf32>
|
|
// CHECK: %[[T0:.+]] = linalg.init_tensor [5, 3, 7]
|
|
// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[T0]]
|
|
// CHECK-SAME: [#[[MAP0]], #[[MAP1]]]
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
|
|
// CHECK-SAME: ins(%[[ARG0]] : tensor<3x5x7xf32>)
|
|
// CHECK-SAME: outs(%[[T1]] : tensor<5x21xf32>)
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
|
|
%arg1 : tensor<?x?x?x5xf32>) ->
|
|
tensor<?x?xf32>
|
|
{
|
|
// expected-remark @+1 {{fused op indexing map is not affine}}
|
|
%0 = linalg.generic {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
|
|
ins(%arg0, %arg1 : tensor<?x?x?x5xf32>, tensor<?x?x?x5xf32>)
|
|
outs(%arg0 : tensor<?x?x?x5xf32>) {
|
|
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // no predecessors
|
|
%1 = mulf %arg3, %arg4 : f32
|
|
linalg.yield %1 : f32
|
|
} -> tensor<?x?x?x5xf32>
|
|
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
|
|
affine_map<(i, j, k, l) -> (j, k, l)>] :
|
|
tensor<?x?x?x5xf32> into tensor<?x?xf32>
|
|
return %1 : tensor<?x?xf32>
|
|
}
|