mlir/test/Dialect/Linalg/inline-scalar-operands.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 @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):  // no predecessors
      // 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 @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):  // no predecessors
      // CHECK: tensor.extract %[[SCALAR]][%[[ZERO]]]
      %2 = arith.divf %arg1, %arg2 : f32
      linalg.yield %2 : f32
    } -> tensor<4xf32>
  return %1 : tensor<4xf32>
}