[Mlir-commits] [mlir] [mlir][quant] Skip non-finite scales in fake quant bounding (PR #209985)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-quant

Author: Jesse Rosenstock (jmr)

<details>
<summary>Changes</summary>

Skip quantization if the dynamically generated scale is not finite. This occurs frequently in quantized JAX models where padding values bounded at -jnp.inf trigger overflow during fake quantization bounding.

Note: FakeQuantSupport.cpp acts as a legacy translation utility whose primary consumers live downstream, and there are currently no in-tree tests covering it. Therefore, no tests were added for this fix.

Assisted-by: Gemini

---
Full diff: https://github.com/llvm/llvm-project/pull/209985.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp (+8) 


``````````diff
diff --git a/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp b/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
index fb27640bfd278..97637ab71ff38 100644
--- a/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
+++ b/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
@@ -131,6 +131,12 @@ mlir::quant::fakeQuantAttrsToType(Location loc, unsigned numBits, double rmin,
   int64_t nudgedZeroPoint;
   getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
 
+  // Skip quantization if the dynamically generated scale is not finite.
+  // This occurs frequently in quantized JAX models where padding values
+  // bounded at -jnp.inf trigger overflow during fake quantization bounding.
+  if (!std::isfinite(scale))
+    return nullptr;
+
   return UniformQuantizedType::getChecked(loc, flags, storageType,
                                           expressedType, scale, nudgedZeroPoint,
                                           qmin, qmax);
@@ -173,6 +179,8 @@ UniformQuantizedPerAxisType mlir::quant::fakeQuantAttrsToType(
     double scale;
     int64_t nudgedZeroPoint;
     getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+    if (!std::isfinite(scale))
+      return nullptr;
     scales.push_back(scale);
     zeroPoints.push_back(nudgedZeroPoint);
   }

``````````

</details>


https://github.com/llvm/llvm-project/pull/209985