mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Mosaic TPU] Optimize clipping impelmentation in arith.fptosi

Browse Source 
We can use maxf/minf to avoid extra comparisons

PiperOrigin-RevId: 720601304
This commit is contained in:
Adam Paszke 2025-01-28 09:19:46 -08:00 committed by jax authors
parent f504d32492
commit 29b658b358
1 changed files with 2 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
View File

@ -597,12 +597,8 @@ LogicalResult canonicalize_fptosi(const CanonicalizeContext &ctx,
}
// Need to clip values to match XLA
auto clip = [&](Value x, Value low, Value high) {
auto is_small =
builder.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, x, low);
x = builder.create<arith::SelectOp>(is_small, low, x);
auto is_large =
builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, x, high);
x = builder.create<arith::SelectOp>(is_large, high, x);
x = builder.create<arith::MaximumFOp>(x, low);
x = builder.create<arith::MinimumFOp>(x, high);
return x;
};
auto minval = builder.getF32FloatAttr(