wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
tensorflow--tensorflow/tensorflow/compiler/mlir/lite/tests/push-tpose-through-ewise.mlir
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

213 lines
11 KiB
MLIR
Raw Permalink Blame History

// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: litert-opt %s --tfl-push-transpose-through-ewise --split-input-file | FileCheck %s
// CHECK-LABEL: pushTposeAfterAddSimple
func.func @pushTposeAfterAddSimple(%arg0: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<5x2x3x4xf32>
%1 = tfl.add %0, %cst { fused_activation_function = "NONE" } : tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %cst_0 = arith.constant dense<1.000000e+00> : tensor<2x3x4x5xf32>
// CHECK: %0 = tfl.add %arg0, %cst_0 {fused_activation_function = "NONE"} : tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: pushTposeAfterAddSimpleWithFold
func.func @pushTposeAfterAddSimpleWithFold(%arg0: tensor<2x3xi32>) -> tensor<3x2xi32> {
%perm = arith.constant dense<[1, 0]> : tensor<2xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3xi32>, tensor<2xi32>) -> tensor<3x2xi32>
%cst = arith.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
%1 = tfl.add %0, %cst { fused_activation_function = "NONE" } : tensor<3x2xi32>
func.return %1 : tensor<3x2xi32>
}
// CHECK: %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
// CHECK: [1, 3, 5], [2, 4, 6]
// CHECK: %0 = tfl.add %arg0, %cst_0 {fused_activation_function = "NONE"} : tensor<2x3xi32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3xi32>, tensor<2xi32>) -> tensor<3x2xi32>
// CHECK: return %1 : tensor<3x2xi32>
// -----
// CHECK-LABEL: pushTposeAfterSubSimple
func.func @pushTposeAfterSubSimple(%arg0: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<5x2x3x4xf32>
%1 = tfl.sub %0, %cst { fused_activation_function = "NONE" } : tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %cst_0 = arith.constant dense<1.000000e+00> : tensor<2x3x4x5xf32>
// CHECK: %0 = tfl.sub %arg0, %cst_0 {fused_activation_function = "NONE"} : tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: permNotConstNoChange
func.func @permNotConstNoChange(%arg0: tensor<2x3x4x5xf32>, %perm: tensor<4xi32>) -> tensor<5x2x3x4xf32> {
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<5x2x3x4xf32>
%1 = tfl.add %0, %cst { fused_activation_function = "NONE" } : tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<1.000000e+00> : tensor<5x2x3x4xf32>
// CHECK: %0 = "tfl.transpose"(%arg0, %arg1) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: %1 = tfl.add %0, %cst {fused_activation_function = "NONE"} : tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: doubleTposeInput
func.func @doubleTposeInput(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%1 = "tfl.transpose"(%arg1, %perm1) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%2 = tfl.add %0, %1 { fused_activation_function = "NONE" } : tensor<5x2x3x4xf32>
func.return %2 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: pushTposeBcastNoChange
func.func @pushTposeBcastNoChange(%arg0: tensor<2x3x4x1xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x1xf32>, tensor<4xi32>) -> tensor<1x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<5x2x3x4xf32>
%1 = "tfl.add"(%0, %cst) { fused_activation_function = "NONE" } : (tensor<1x2x3x4xf32>, tensor<5x2x3x4xf32>) -> tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<1.000000e+00> : tensor<5x2x3x4xf32>
// CHECK: %cst_0 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst_0) : (tensor<2x3x4x1xf32>, tensor<4xi32>) -> tensor<1x2x3x4xf32>
// CHECK: %1 = tfl.add(%0, %cst) <{fused_activation_function = "NONE"}> : (tensor<1x2x3x4xf32>, tensor<5x2x3x4xf32>) -> tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: doubleTposeOneBroadcastInput
func.func @doubleTposeOneBroadcastInput(%arg0: tensor<2x3x4x1xf32>, %arg1: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x1xf32>, tensor<4xi32>) -> tensor<1x2x3x4xf32>
%perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%1 = "tfl.transpose"(%arg1, %perm1) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%2 = "tfl.add"(%0, %1) { fused_activation_function = "NONE" } : (tensor<1x2x3x4xf32>, tensor<5x2x3x4xf32>) -> tensor<5x2x3x4xf32>
func.return %2 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = tfl.add(%arg0, %arg1) <{fused_activation_function = "NONE"}> : (tensor<2x3x4x1xf32>, tensor<2x3x4x5xf32>) -> tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: doubleTposeInputPermNotEqualNoChange
func.func @doubleTposeInputPermNotEqualNoChange(%arg0: tensor<2x4x3x5xf32>, %arg1: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 2, 1]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x4x3x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%1 = "tfl.transpose"(%arg1, %perm1) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%2 = tfl.add %0, %1 { fused_activation_function = "NONE" } : tensor<5x2x3x4xf32>
func.return %2 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %cst_0 = arith.constant dense<[3, 0, 2, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst_0) : (tensor<2x4x3x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: %1 = "tfl.transpose"(%arg1, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: %2 = tfl.add %0, %1 {fused_activation_function = "NONE"} : tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: pushTposeBcastCstInput
func.func @pushTposeBcastCstInput(%arg0: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<1x2x3x4xf32>
%1 = "tfl.add"(%0, %cst) { fused_activation_function = "NONE" } : (tensor<5x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %cst_0 = arith.constant dense<1.000000e+00> : tensor<2x3x4x1xf32>
// CHECK: %0 = tfl.add(%arg0, %cst_0) <{fused_activation_function = "NONE"}> : (tensor<2x3x4x5xf32>, tensor<2x3x4x1xf32>) -> tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: pushTposeBcastScalarCstInput
func.func @pushTposeBcastScalarCstInput(%arg0: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
%cst = arith.constant dense<1.0> : tensor<f32>
%1 = "tfl.add"(%0, %cst) { fused_activation_function = "NONE" } : (tensor<5x2x3x4xf32>, tensor<f32>) -> tensor<5x2x3x4xf32>
func.return %1 : tensor<5x2x3x4xf32>
}
// CHECK: %cst = arith.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %cst_0 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = tfl.add(%arg0, %cst) <{fused_activation_function = "NONE"}> : (tensor<2x3x4x5xf32>, tensor<f32>) -> tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst_0) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// -----
// CHECK-LABEL: pushTposeDynamicBcastScalarCstInput
func.func @pushTposeDynamicBcastScalarCstInput(%arg0: tensor<?x?x4x5xf32>) -> tensor<5x?x?x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<?x?x4x5xf32>, tensor<4xi32>) -> tensor<5x?x?x4xf32>
%cst = arith.constant dense<1.0> : tensor<f32>
%1 = "tfl.add"(%0, %cst) { fused_activation_function = "NONE" } : (tensor<5x?x?x4xf32>, tensor<f32>) -> tensor<5x?x?x4xf32>
func.return %1 : tensor<5x?x?x4xf32>
}
// CHECK: %cst = arith.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %cst_0 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = tfl.add(%arg0, %cst) <{fused_activation_function = "NONE"}> : (tensor<?x?x4x5xf32>, tensor<f32>) -> tensor<?x?x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst_0) : (tensor<?x?x4x5xf32>, tensor<4xi32>) -> tensor<5x?x?x4xf32>
// -----
// CHECK-LABEL: doubleTposeDynamicInput
func.func @doubleTposeDynamicInput(%arg0: tensor<?x?x4x5xf32>, %arg1: tensor<?x?x4x5xf32>) -> tensor<5x?x?x4xf32> {
%perm = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%0 = "tfl.transpose"(%arg0, %perm) : (tensor<?x?x4x5xf32>, tensor<4xi32>) -> tensor<5x?x?x4xf32>
%perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
%1 = "tfl.transpose"(%arg1, %perm1) : (tensor<?x?x4x5xf32>, tensor<4xi32>) -> tensor<5x?x?x4xf32>
%2 = tfl.add %0, %1 { fused_activation_function = "NONE" } : tensor<5x?x?x4xf32>
func.return %2 : tensor<5x?x?x4xf32>
}
// CHECK: %cst = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>
// CHECK: %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<?x?x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<?x?x4x5xf32>, tensor<4xi32>) -> tensor<5x?x?x4xf32>
// CHECK: return %1 : tensor<5x?x?x4xf32>
Reference in New Issue View Git Blame Copy Permalink