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

[MOSAIC] Accept kernels for a compatibility window with MAXNUMF instead of MAXIMUMF

Browse Source 
PiperOrigin-RevId: 605436599
This commit is contained in:
Blake Hechtman 2024-02-08 14:56:47 -08:00 committed by jax authors
parent 8139e58e02
commit d29c86eb52
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
View File

@ -2437,10 +2437,12 @@ LogicalResult vector_multi_reduction_rule(RewriteContext &ctx, Operation &op,
case vector::CombiningKind::ADD:
neutral = builder.getF32FloatAttr(0);
break;
case vector::CombiningKind::MAXNUMF:
case vector::CombiningKind::MAXIMUMF: {
// TODO(b/322836633): The semantics of maximumf don't match the lowering
// for older TPU versions because older TPU versions don't respect the
// -0.0 vs +0.0 ordering.
// -0.0 vs +0.0 ordering. Keeping MAXNUMF for backward compatibility of
// serialized artifacts.
neutral = builder.getFloatAttr(
builder.getF32Type(),
APFloat::getInf(APFloat::IEEEsingle(), /*Negative=*/true));
@ -2532,6 +2534,7 @@ LogicalResult vector_multi_reduction_rule(RewriteContext &ctx, Operation &op,
case vector::CombiningKind::ADD:
tpu_kind = tpu::ReductionKind::SUM;
break;
case vector::CombiningKind::MAXNUMF:
case vector::CombiningKind::MAXIMUMF:
tpu_kind = tpu::ReductionKind::MAX;
break;