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

[Mosaic] Remove "support" for MAXNUMF vector reductions

Browse Source 
Pallas hasn't been using it to lower since cl/604849222, which is before we had serde, so we won't create a serde rule for this. That CL also tried to remove the support, but it had to be restored in cl/605436599 because it was breaking serialized kernels.

PiperOrigin-RevId: 639896981
This commit is contained in:
Tomás Longeri 2024-06-03 13:39:47 -07:00 committed by jax authors
parent 2bfafe35d2
commit e620acfa17
1 changed files with 1 additions and 4 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

@ -3319,12 +3319,10 @@ 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. Keeping MAXNUMF for backward compatibility of
// serialized artifacts.
// -0.0 vs +0.0 ordering.
neutral = builder.getFloatAttr(
builder.getF32Type(),
APFloat::getInf(APFloat::IEEEsingle(), /*Negative=*/true));
@ -3418,7 +3416,6 @@ 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;