diff --git a/mlir/test/Dialect/Tensor/ops.mlir b/mlir/test/Dialect/Tensor/ops.mlir --- a/mlir/test/Dialect/Tensor/ops.mlir +++ b/mlir/test/Dialect/Tensor/ops.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt <%s | mlir-opt | FileCheck %s +// RUN: mlir-opt %s | mlir-opt | FileCheck %s // CHECK-LABEL: func @cast( func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor) { @@ -29,7 +29,6 @@ return %0 : tensor<5x?x6xf32, "foo"> } - // CHECK-LABEL: func @extract( // CHECK-SAME: %[[TENSOR:.*]]: tensor, // CHECK-SAME: %[[INDEX:.*]]: index) { @@ -121,8 +120,6 @@ return } -// ----- - // CHECK-LABEL: func @insert_slice({{.*}}) { func.func @insert_slice( %t: tensor<8x16x4xf32>, @@ -157,8 +154,6 @@ return } -// ----- - func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor) -> (tensor, tensor<1x1xf32>) { %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor @@ -180,8 +175,6 @@ // CHECK: tensor.collapse_shape // CHECK-SAME: [0], [1], [2, 3, 4] -// ----- - func.func @rank(%t : tensor<4x4x?xf32>) { // CHECK: %{{.*}} = tensor.rank %{{.*}} : tensor<4x4x?xf32> %0 = "tensor.rank"(%t) : (tensor<4x4x?xf32>) -> index @@ -191,8 +184,6 @@ return } -// ----- - func.func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, %pad_value: f32) -> tensor<6x?x?x?xf32> { %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] { @@ -210,8 +201,6 @@ // CHECK-SAME: high[3, 3, %[[HIGH]], 2] // CHECK: : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32> -// ----- - func.func @pad_static(%arg0: tensor<3x4xf32>, %pad_value: f32) -> tensor<6x9xf32> { %0 = tensor.pad %arg0 low[1, 2] high[2, 3] { ^bb0(%arg1 : index, %arg2 : index): @@ -224,8 +213,6 @@ // CHECK: tensor.pad %[[ARG0]] low[1, 2] high[2, 3] // CHECK: : tensor<3x4xf32> to tensor<6x9xf32> -// ----- - func.func @pad_asymmetrical(%arg0: tensor<2x3xf32>, %ub0: index, %ub1: index, %pad_value: f32) -> tensor { %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] { @@ -243,8 +230,6 @@ // CHECK-SAME: high[%[[UB0]], %[[UB1]]] // CHECK: : tensor<2x3xf32> to tensor -// ----- - func.func @pad_to_static_size(%arg0: tensor, %ub0: index, %ub1: index, %pad_value: f32) -> tensor<2x3xf32> { %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] { @@ -262,8 +247,6 @@ // CHECK-SAME: high[%[[UB0]], %[[UB1]]] // CHECK: : tensor to tensor<2x3xf32> -// ----- - // CHECK-LABEL: func @test_splat_op // CHECK-SAME: [[S:%arg[0-9]+]]: f32 func.func @test_splat_op(%s : f32) { @@ -275,19 +258,24 @@ return } -// ----- - -// CHECK-LABEL: func @gather_scatter +// CHECK-LABEL: func.func @gather_scatter( +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5x6xf32>, +// CHECK-SAME: %[[ARG1:.*]]: tensor<1x3x2xindex>, +// CHECK-SAME: %[[ARG2:.*]]: tensor<1x3x2xi32>) { func.func @gather_scatter( %dest : tensor<4x5x6xf32>, %indices: tensor<1x3x2xindex>, %indices_i32: tensor<1x3x2xi32>) { + // CHECK: %[[GATHER:.*]] = tensor.gather %[[ARG0]][%[[ARG2]]] gather_dims([1, 2]) unique : (tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<1x3x4x1x1xf32> %gathered = tensor.gather %dest[%indices_i32] gather_dims([1, 2]) unique: (tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<1x3x4x1x1xf32> + // CHECK: %[[GATHER0:.*]] = tensor.gather %[[ARG0]][%[[ARG1]]] gather_dims([1, 2]) unique : (tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<1x3x4xf32> %rank_reduced_gathered = tensor.gather %dest[%indices] gather_dims([1, 2]) unique: (tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<1x3x4xf32> + // CHECK: %{{.*}} = tensor.scatter %[[GATHER]] into %[[ARG0]][%[[ARG1]]] scatter_dims([1, 2]) unique : (tensor<1x3x4x1x1xf32>, tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<4x5x6xf32> %scattered = tensor.scatter %gathered into %dest[%indices] scatter_dims([1, 2]) unique: (tensor<1x3x4x1x1xf32>, tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<4x5x6xf32> + // CHECK: %{{.*}} = tensor.scatter %[[GATHER0]] into %[[ARG0]][%[[ARG2]]] scatter_dims([1, 2]) unique : (tensor<1x3x4xf32>, tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<4x5x6xf32> %rank_reduced_scattered = tensor.scatter %rank_reduced_gathered into %dest[%indices_i32] scatter_dims([1, 2]) unique: (tensor<1x3x4xf32>, tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<4x5x6xf32>