27 lines
1.4 KiB
MLIR
27 lines
1.4 KiB
MLIR
// RUN: mlir-hlo-opt --chlo-legalize-to-hlo --split-input-file %s | FileCheck %s
|
|
|
|
// Lower statically shaped `constant_like` to constant.
|
|
// CHECK-LABEL: @constant_like_static_shape
|
|
func @constant_like_static_shape(%arg : tensor<1x2xi64>) -> tensor<1x2xf32> {
|
|
// CHECK: %[[RESULT:.*]] = mhlo.constant dense<3.200000e+00> : tensor<1x2xf32>
|
|
// CHECK: return %[[RESULT]]
|
|
%result = "chlo.constant_like"(%arg) { value = 3.2 : f32 }
|
|
: (tensor<1x2xi64>) -> tensor<1x2xf32>
|
|
return %result : tensor<1x2xf32>
|
|
}
|
|
|
|
// Lower dynamically shaped `constant_like` to broadcasted constant.
|
|
// CHECK-LABEL: constant_like_dynamic_shape
|
|
// CHECK-SAME: (%[[ARG:.*]]: tensor<?x?xi64>)
|
|
func @constant_like_dynamic_shape(%arg : tensor<?x?xi64>) -> tensor<?x?xf32> {
|
|
// CHECK: %[[CONSTANT:.*]] = mhlo.constant dense<3.200000e+00> : tensor<f32>
|
|
// CHECK: %[[UNCASTED_SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<?x?xi64> -> tensor<?xindex>
|
|
// CHECK: %[[SHAPE:.*]] = tensor_cast %[[UNCASTED_SHAPE]] : tensor<?xindex> to tensor<2xindex>
|
|
// CHECK: %[[BROADCASTED_CONSTANT:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[CONSTANT]], %[[SHAPE]]) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x?xf32>
|
|
// CHECK: return %[[BROADCASTED_CONSTANT]] : tensor<?x?xf32>
|
|
%result = "chlo.constant_like"(%arg) { value = 3.2 : f32 }
|
|
: (tensor<?x?xi64>) -> tensor<?x?xf32>
|
|
return %result : tensor<?x?xf32>
|
|
}
|
|
|