colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[MLIR][KernelGen] Lower mhlo.log_plus_one to std.log1p 

PiperOrigin-RevId: 353200069
This commit is contained in:
A. Unique TensorFlower 2021-01-22 02:16:26 -08:00 committed by TensorFlow MLIR Team
parent f6bf9d5780
commit 56758a9562
2 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h
View File

@ -442,11 +442,8 @@ inline Value MapLhloOpToStdScalarOp<lmhlo::Log1pOp>(Location loc,
ArrayRef<Type> result_types,
ArrayRef<Value> args,
OpBuilder* b) {
auto ty = result_types.front().cast<FloatType>();
Value x = args.front();
Value one = b->create<ConstantOp>(loc, b->getFloatAttr(ty, 1.0));
Value x_plus_one = b->create<AddFOp>(loc, x, one);
return b->create<::mlir::LogOp>(loc, x_plus_one);
return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::Log1pOp>{}(
loc, result_types, args, b);
}
template <>

10
tests/hlo-legalize-to-linalg.mlir
View File

@ -166,6 +166,16 @@ func @float_log(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
// -----
// CHECK-LABEL: func @float_log1p
func @float_log1p(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
// CHECK: linalg.generic
// CHECK: log1p
%0 = "mhlo.log_plus_one"(%arg0) : (tensor<2x2xf32>) -> tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
// -----
// CHECK-LABEL: func @float_ceil
func @float_ceil(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
// CHECK: linalg.generic