From f75d164eeaf407b21ec6b2cff4bcc3ad2003af61 Mon Sep 17 00:00:00 2001 From: srcarroll <50210727+srcarroll@users.noreply.github.com> Date: Fri, 15 Mar 2024 14:39:57 -0500 Subject: [PATCH] =?UTF-8?q?Revert=20"[mlir][math]=20Implement=20alternativ?= =?UTF-8?q?e=20decomposition=20for=20tanh=20(#8=E2=80=A6=20(#85429)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …5025)" This reverts commit 58ef9bec071383744fb703ff08df9806f25e4095. There is a bool to float casting issue that needs to be sorted out to make sure this is target independent --- .../Math/Transforms/ExpandPatterns.cpp | 40 ++++++++----------- mlir/test/Dialect/Math/expand-math.mlir | 19 ++++----- 2 files changed, 27 insertions(+), 32 deletions(-) diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp index 1750171b81a10..989a3e5536ec6 100644 --- a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp +++ b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp @@ -91,40 +91,34 @@ static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) { } /// Expands tanh op into -/// 1-exp^{-2x} / 1+exp^{-2x} -/// To avoid overflow we exploit the reflection symmetry `tanh(-x) = -tanh(x)`. -/// We compute a "signs" value which is -1 if input is negative and +1 if input -/// is positive. Then multiply the input by this value, guaranteeing that the -/// result is positive, which also guarantees `exp^{-2x * sign(x)}` is in (0, -/// 1]. Expand the computation on the input `x * sign(x)`, then multiply the -/// result by `sign(x)` to retain sign of the real result. +/// 1) 1-exp^{-2x} / 1+exp^{-2x}, if x => 0 +/// 2) exp^{2x}-1 / exp^{2x}+1 , if x < 0 static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) { auto floatType = op.getOperand().getType(); Location loc = op.getLoc(); - Value zero = createFloatConst(loc, floatType, 0.0, rewriter); Value one = createFloatConst(loc, floatType, 1.0, rewriter); - Value negTwo = createFloatConst(loc, floatType, -2.0, rewriter); - - // Compute sign(x) = cast(x < 0) * (-2) + 1 - Value sign = rewriter.create(loc, arith::CmpFPredicate::OLT, - op.getOperand(), zero); - sign = rewriter.create(loc, floatType, sign); - sign = rewriter.create(loc, sign, negTwo); - sign = rewriter.create(loc, sign, one); + Value two = createFloatConst(loc, floatType, 2.0, rewriter); + Value doubledX = rewriter.create(loc, op.getOperand(), two); - // Normalize input to positive value: y = sign(x) * x - Value positiveX = rewriter.create(loc, sign, op.getOperand()); - - // Decompose on normalized input - Value negDoubledX = rewriter.create(loc, negTwo, positiveX); + // Case 1: tanh(x) = 1-exp^{-2x} / 1+exp^{-2x} + Value negDoubledX = rewriter.create(loc, doubledX); Value exp2x = rewriter.create(loc, negDoubledX); Value dividend = rewriter.create(loc, one, exp2x); Value divisor = rewriter.create(loc, one, exp2x); Value positiveRes = rewriter.create(loc, dividend, divisor); - // Multiply result by sign(x) to retain signs from negative inputs - rewriter.replaceOpWithNewOp(op, sign, positiveRes); + // Case 2: tanh(x) = exp^{2x}-1 / exp^{2x}+1 + exp2x = rewriter.create(loc, doubledX); + dividend = rewriter.create(loc, exp2x, one); + divisor = rewriter.create(loc, exp2x, one); + Value negativeRes = rewriter.create(loc, dividend, divisor); + // tanh(x) = x >= 0 ? positiveRes : negativeRes + Value zero = createFloatConst(loc, floatType, 0.0, rewriter); + Value cmpRes = rewriter.create(loc, arith::CmpFPredicate::OGE, + op.getOperand(), zero); + rewriter.replaceOpWithNewOp(op, cmpRes, positiveRes, + negativeRes); return success(); } diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir index 86ee5c8620472..6ee65b085dad1 100644 --- a/mlir/test/Dialect/Math/expand-math.mlir +++ b/mlir/test/Dialect/Math/expand-math.mlir @@ -7,18 +7,19 @@ func.func @tanh(%arg: f32) -> f32 { } // CHECK-DAG: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ONE:.+]] = arith.constant 1.000000e+00 : f32 -// CHECK-DAG: %[[TWO:.+]] = arith.constant -2.000000e+00 : f32 -// CHECK: %[[VAL0:.+]] = arith.cmpf olt, %arg0, %[[ZERO]] : f32 -// CHECK: %[[VAL1:.+]] = arith.sitofp %[[VAL0]] : i1 to f32 -// CHECK: %[[VAL2:.+]] = arith.mulf %[[VAL1]], %[[TWO]] : f32 -// CHECK: %[[SIGN:.+]] = arith.addf %[[VAL2]], %[[ONE]] : f32 -// CHECK: %[[POSX:.+]] = arith.mulf %[[SIGN]], %arg0 : f32 -// CHECK: %[[NEGDOUBLEDX:.+]] = arith.mulf %[[POSX]], %[[TWO]] : f32 +// CHECK-DAG: %[[TWO:.+]] = arith.constant 2.000000e+00 : f32 +// CHECK: %[[DOUBLEDX:.+]] = arith.mulf %arg0, %[[TWO]] : f32 +// CHECK: %[[NEGDOUBLEDX:.+]] = arith.negf %[[DOUBLEDX]] : f32 // CHECK: %[[EXP1:.+]] = math.exp %[[NEGDOUBLEDX]] : f32 // CHECK: %[[DIVIDEND1:.+]] = arith.subf %[[ONE]], %[[EXP1]] : f32 // CHECK: %[[DIVISOR1:.+]] = arith.addf %[[EXP1]], %[[ONE]] : f32 -// CHECK: %[[POSRES:.+]] = arith.divf %[[DIVIDEND1]], %[[DIVISOR1]] : f32 -// CHECK: %[[RESULT:.+]] = arith.mulf %[[SIGN]], %[[POSRES]] : f32 +// CHECK: %[[RES1:.+]] = arith.divf %[[DIVIDEND1]], %[[DIVISOR1]] : f32 +// CHECK: %[[EXP2:.+]] = math.exp %[[DOUBLEDX]] : f32 +// CHECK: %[[DIVIDEND2:.+]] = arith.subf %[[EXP2]], %[[ONE]] : f32 +// CHECK: %[[DIVISOR2:.+]] = arith.addf %[[EXP2]], %[[ONE]] : f32 +// CHECK: %[[RES2:.+]] = arith.divf %[[DIVIDEND2]], %[[DIVISOR2]] : f32 +// CHECK: %[[COND:.+]] = arith.cmpf oge, %arg0, %[[ZERO]] : f32 +// CHECK: %[[RESULT:.+]] = arith.select %[[COND]], %[[RES1]], %[[RES2]] : f32 // CHECK: return %[[RESULT]] // -----