forked from OSchip/llvm-project
49 lines
2.1 KiB
MLIR
49 lines
2.1 KiB
MLIR
// RUN: mlir-opt %s -linalg-inline-scalar-operands -split-input-file | FileCheck %s
|
|
|
|
// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>
|
|
#map2 = affine_map<(d0) -> (d0)>
|
|
#map3 = affine_map<(d0) -> ()>
|
|
|
|
// CHECK: func @inline_zerod(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<f32>)
|
|
func.func @inline_zerod(%arg0: tensor<4xf32>, %scalar: tensor<f32>) -> tensor<4xf32> {
|
|
%0 = linalg.init_tensor [4] : tensor<4xf32>
|
|
// CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],
|
|
// CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)
|
|
%1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],
|
|
iterator_types = ["parallel"]}
|
|
ins(%arg0, %scalar : tensor<4xf32>, tensor<f32>)
|
|
outs(%0 : tensor<4xf32>) {
|
|
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)
|
|
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
|
|
// CHECK: tensor.extract %[[SCALAR]][]
|
|
%2 = arith.divf %arg1, %arg2 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<4xf32>
|
|
return %1 : tensor<4xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>
|
|
#map2 = affine_map<(d0) -> (d0)>
|
|
#map3 = affine_map<(d0) -> (0)>
|
|
|
|
// CHECK: func @inline_oned(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<1xf32>)
|
|
func.func @inline_oned(%arg0: tensor<4xf32>, %scalar: tensor<1xf32>) -> tensor<4xf32> {
|
|
// CHECK: %[[ZERO:.*]] = arith.constant 0 : index
|
|
%0 = linalg.init_tensor [4] : tensor<4xf32>
|
|
// CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],
|
|
// CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)
|
|
%1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],
|
|
iterator_types = ["parallel"]}
|
|
ins(%arg0, %scalar : tensor<4xf32>, tensor<1xf32>)
|
|
outs(%0 : tensor<4xf32>) {
|
|
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)
|
|
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
|
|
// CHECK: tensor.extract %[[SCALAR]][%[[ZERO]]]
|
|
%2 = arith.divf %arg1, %arg2 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<4xf32>
|
|
return %1 : tensor<4xf32>
|
|
}
|