[Torch] Fix aten.linalg_vector_norm for ord = 0 and ord = +/-inf (#4718)

Fix `aten.linalg_vector_norm` for `ord = 0` and `ord = ±inf`

The generic `(sum |x|^ord)^(1/ord)` lowering is wrong for these:
`±inf` gives `inf^0 = 1` and `ord = 0` gives `pow(N, +inf)`, so
`ord = -inf` was returning `1.0` instead of `min(|x_i|)`.

Instead of patching each backend, I decompose the three cases in
`DecomposeComplexOps` (following `_refs`): `+inf → amax(|x|)`,
`-inf → amin(|x|)`, `0 → sum(|x| != 0)`. This also fixes the ONNX path
(`LpNormalization → aten.norm.ScalarOpt_dim → aten.linalg_vector_norm`),
which never hits the Python `_refs` decomposition. Finite / non-constant
`ord` still go through the backends unchanged, which now decline
`0`/`±inf` cleanly instead of miscompiling.

Also drops the StableHLO XFAILs — the decomposition makes these correct
there too.

Author: hariprasadravi

lib/Conversion/TorchToLinalg/Reduction.cpp

@@ -678,16 +678,6 @@ class ConvertReductionOp : public ConversionPattern {
     // Cast `ord` to float so that we can readily pass it math.powf.
     Value ordValue = convertScalarToDtype(rewriter, loc, ordOp, elemType);
 
-    // TODO: Add support for ord = {0, +inf, -inf}.
-    auto epsilon = 1e-5;
-    auto ordLiteral = 0.0;
-    if (matchPattern(ordValue, m_TorchConstantFloat(&ordLiteral)) &&
-        fabs(ordLiteral) < epsilon)
-      return rewriter.notifyMatchFailure(op, "unimplemented: L0 norm");
-
-    if (std::isinf(ordLiteral))
-      return rewriter.notifyMatchFailure(op, "unimplemented: ord = +/- inf");
-
     // Raise each summed value to the inverse of the order of the norm.
     TypedAttr oneAttr = rewriter.getFloatAttr(elemType, 1.0);
     auto oneValue = arith::ConstantOp::create(rewriter, loc, oneAttr);
@@ -759,6 +749,26 @@ class ConvertReductionOp : public ConversionPattern {
       return rewriter.notifyMatchFailure(
           op, "invalid operand or result types to use with linalg on tensors");
 
+    // ord = 0 (count of nonzeros, imported as an int) and ord = +/-inf (min/max
+    // of absolute values, imported as a float) are handled by
+    // DecomposeAtenLinalgVectorNormOp; the generic (sum |x|^ord)^(1/ord)
+    // lowering is undefined for them. Decline before creating any IR so that a
+    // miscompile cannot slip through if decomposition is disabled. Match on the
+    // original `ord` scalar; a non-constant ord is left to the generic path.
+    if (auto normOp = dyn_cast<AtenLinalgVectorNormOp>(op)) {
+      double ordLiteral;
+      int64_t ordInt;
+      bool isConstOrd = true;
+      if (matchPattern(normOp.getOrd(), m_TorchConstantInt(&ordInt)))
+        ordLiteral = static_cast<double>(ordInt);
+      else if (!matchPattern(normOp.getOrd(),
+                             m_TorchConstantFloat(&ordLiteral)))
+        isConstOrd = false;
+      if (isConstOrd && (ordLiteral == 0.0 || std::isinf(ordLiteral)))
+        return rewriter.notifyMatchFailure(
+            op, "ord = 0 / +/-inf are handled by decomposition");
+    }
+
     FailureOr<torch_to_linalg::ReductionOpInfo> opInfo =
         computeReductionOpInfo(op, operands, rewriter);
     if (failed(opInfo))

lib/Conversion/TorchToStablehlo/Reduction.cpp

@@ -910,6 +910,23 @@ LogicalResult ConvertAtenReductionOp<AtenLinalgVectorNormOp>::matchAndRewrite(
   }
   int64_t inputRank = inputType.getRank();
 
+  // ord = 0 (count of nonzeros, imported as an int) and ord = +/-inf (min/max
+  // of absolute values, imported as a float) are handled by
+  // DecomposeAtenLinalgVectorNormOp; the generic (sum |x|^ord)^(1/ord) lowering
+  // below is undefined for them. Decline before creating any IR so that a
+  // miscompile cannot slip through if decomposition is disabled. Match on the
+  // original `ord` scalar; a non-constant ord is left to the generic path.
+  double ordFloat;
+  int64_t ordInt;
+  bool isConstOrd = true;
+  if (matchPattern(op.getOrd(), m_TorchConstantInt(&ordInt)))
+    ordFloat = static_cast<double>(ordInt);
+  else if (!matchPattern(op.getOrd(), m_TorchConstantFloat(&ordFloat)))
+    isConstOrd = false;
+  if (isConstOrd && (ordFloat == 0.0 || std::isinf(ordFloat)))
+    return rewriter.notifyMatchFailure(
+        op, "ord = 0 / +/-inf are handled by decomposition");
+
   auto outType =
       cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
   auto outElemType = outType.getElementType();

lib/Conversion/TorchToTosa/TosaLegalizeCommon.cpp

@@ -1121,44 +1121,46 @@ convertLinalgVectorNormOp(PatternRewriter &rewriter, Operation *op,
   }
 
   auto linalgVectorNormOp = cast<AtenLinalgVectorNormOp>(op);
-  // TODO: Add support for ord = {0, +inf, -inf}.
-  auto epsilon = 1e-5;
   double ordLiteralFloat = 1.0;
   int64_t ordLiteralInt = 1;
-  Value ordVal;
-  if (matchPattern(linalgVectorNormOp.getOrd(),
-                   torch::Torch::m_TorchConstantFloat(&ordLiteralFloat))) {
-    ordVal = tosa::getConstTensor<float>(rewriter, op,
-                                         {static_cast<float>(ordLiteralFloat)},
-                                         {}, elemType)
-                 .value();
-  } else if (matchPattern(linalgVectorNormOp.getOrd(),
-                          torch::Torch::m_TorchConstantInt(&ordLiteralInt))) {
-    ordVal = tosa::getConstTensor<float>(rewriter, op,
-                                         {static_cast<float>(ordLiteralInt)},
-                                         {}, elemType)
-                 .value();
-  } else {
+  bool ordIsFloat =
+      matchPattern(linalgVectorNormOp.getOrd(),
+                   torch::Torch::m_TorchConstantFloat(&ordLiteralFloat));
+  bool ordIsInt =
+      !ordIsFloat &&
+      matchPattern(linalgVectorNormOp.getOrd(),
+                   torch::Torch::m_TorchConstantInt(&ordLiteralInt));
+  if (!ordIsFloat && !ordIsInt) {
     op->emitOpError("only support FP or INT type ord parameter");
     return std::nullopt;
   }
 
-  Value ordValRank0 = ordVal;
-  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), input_value, ordVal)
-          .failed())
+  // ord = 0 (count of nonzeros) and ord = +/-inf (min/max of absolute values)
+  // are handled by DecomposeAtenLinalgVectorNormOp; the generic
+  // (sum |x|^ord)^(1/ord) lowering below is undefined for them. Decline before
+  // creating any IR so that a miscompile cannot slip through if decomposition
+  // is disabled.
+  double ordLiteral =
+      ordIsFloat ? ordLiteralFloat : static_cast<double>(ordLiteralInt);
+  if (ordLiteral == 0.0) {
+    (void)rewriter.notifyMatchFailure(op,
+                                      "ord = 0 is handled by decomposition");
     return std::nullopt;
-
-  if (fabs(ordLiteralFloat) < epsilon ||
-      fabs(static_cast<double>(ordLiteralInt)) < epsilon) {
-    op->emitOpError("unimplemented: L0 norm");
+  }
+  if (std::isinf(ordLiteral)) {
+    (void)rewriter.notifyMatchFailure(
+        op, "ord = +/-inf are handled by decomposition");
     return std::nullopt;
   }
 
-  if (std::isinf(ordLiteralFloat) ||
-      std::isinf(static_cast<double>(ordLiteralInt))) {
-    op->emitOpError("unimplemented: ord = +/- inf");
+  Value ordVal = tosa::getConstTensor<float>(rewriter, op,
+                                             {static_cast<float>(ordLiteral)},
+                                             {}, elemType)
+                     .value();
+  Value ordValRank0 = ordVal;
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), input_value, ordVal)
+          .failed())
     return std::nullopt;
-  }
 
   auto input_value_casted =
       tosa::tosaCastTensorToType(rewriter, input_value, output_type).value();

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -10677,6 +10677,82 @@ class DecomposeAtenNormScalarOptDimOp
 };
 } // namespace
 
+namespace {
+// Decompose `aten.linalg_vector_norm` for the `ord` values where the generic
+// `(sum |x|^ord)^(1/ord)` lowering is undefined: `ord = +inf` is `max(|x|)`,
+// `ord = -inf` is `min(|x|)`, and `ord = 0` is the count of nonzero elements
+// `sum(|x| != 0)`. Handling these here makes every backend correct through the
+// already-supported `aten.amax`/`aten.amin`/`aten.sum.dim_IntList` ops. Finite,
+// nonzero `ord` (and non-constant `ord`) are left for the backends' generic
+// lowering.
+class DecomposeAtenLinalgVectorNormOp
+    : public OpRewritePattern<AtenLinalgVectorNormOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenLinalgVectorNormOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+
+    // `ord` is an `AnyTorchScalarType`, so an integer `ord` (e.g. `ord = 0`)
+    // imports as a `torch.constant.int`, not a `torch.constant.float`.
+    double ordLiteral;
+    int64_t ordInt;
+    if (matchPattern(op.getOrd(), m_TorchConstantInt(&ordInt)))
+      ordLiteral = static_cast<double>(ordInt);
+    else if (!matchPattern(op.getOrd(), m_TorchConstantFloat(&ordLiteral)))
+      return rewriter.notifyMatchFailure(op, "non-constant `ord` unsupported");
+
+    bool isInf = std::isinf(ordLiteral);
+    bool isZero = ordLiteral == 0.0;
+    if (!isInf && !isZero)
+      return rewriter.notifyMatchFailure(
+          op, "only ord = 0 / +/-inf are decomposed; finite ord is lowered by "
+              "the backends");
+
+    Value self = op.getSelf();
+    auto selfType = dyn_cast<BaseTensorType>(self.getType());
+    if (!selfType || !selfType.hasSizes())
+      return rewriter.notifyMatchFailure(op, "expected input with known rank");
+
+    // `dim = None` means reduce over all dimensions.
+    Value dim = op.getDim();
+    if (isa<Torch::NoneType>(dim.getType())) {
+      SmallVector<Value> allDims;
+      for (int64_t i = 0, rank = selfType.getSizes().size(); i < rank; ++i)
+        allDims.push_back(ConstantIntOp::create(rewriter, loc,
+                                                rewriter.getI64IntegerAttr(i)));
+      dim = PrimListConstructOp::create(
+          rewriter, loc,
+          Torch::ListType::get(Torch::IntType::get(op.getContext())), allDims);
+    }
+
+    Value abs = AtenAbsOp::create(rewriter, loc, self.getType(), self);
+
+    // `ord = +/-inf` are the max/min of the absolute values.
+    if (isInf) {
+      if (ordLiteral > 0)
+        rewriter.replaceOpWithNewOp<AtenAmaxOp>(op, op.getType(), abs, dim,
+                                                op.getKeepdim());
+      else
+        rewriter.replaceOpWithNewOp<AtenAminOp>(op, op.getType(), abs, dim,
+                                                op.getKeepdim());
+      return success();
+    }
+
+    // `ord = 0` is the count of nonzero elements: `sum(|x| != 0)`.
+    Value zero =
+        ConstantFloatOp::create(rewriter, loc, rewriter.getF64FloatAttr(0.0));
+    auto boolType = selfType.getWithSizesAndDtype(selfType.getOptionalSizes(),
+                                                  rewriter.getI1Type());
+    Value nonzero = AtenNeScalarOp::create(rewriter, loc, boolType, abs, zero);
+    Value none = ConstantNoneOp::create(rewriter, loc);
+    rewriter.replaceOpWithNewOp<AtenSumDimIntListOp>(
+        op, op.getType(), nonzero, dim, op.getKeepdim(), /*dtype=*/none);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeAtenRandintLowOp : public OpRewritePattern<AtenRandintLowOp> {
 public:
@@ -13842,6 +13918,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenMseLossOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenL1LossOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenNormScalarOptDimOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenLinalgVectorNormOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenRandintOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenRandintLowOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenVarMeanCorrectionOp>(patterns);

projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

@@ -2069,6 +2069,149 @@ def ReduceLN3NormModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ReduceLInfNormModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=0, ord=float("inf"))
+
+
+@register_test_case(module_factory=lambda: ReduceLInfNormModule())
+def ReduceLInfNormModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5))
+
+
+# ==============================================================================
+
+
+class ReduceLNegInfNormModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=0, ord=float("-inf"))
+
+
+@register_test_case(module_factory=lambda: ReduceLNegInfNormModule())
+def ReduceLNegInfNormModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5))
+
+
+# ==============================================================================
+
+
+class ReduceL0NormModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=0, ord=0)
+
+
+@register_test_case(module_factory=lambda: ReduceL0NormModule())
+def ReduceL0NormModule_basic(module, tu: TestUtils):
+    # Include exact zeros so the nonzero-count semantics are exercised.
+    module.forward(torch.tensor([0.0, 1.5, 0.0, -2.0, 3.0]))
+
+
+# ==============================================================================
+
+
+# ord = +inf over a single dim of a multidim input, with keepdim.
+class ReduceLInfNormKeepDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=1, keepdim=True, ord=float("inf"))
+
+
+@register_test_case(module_factory=lambda: ReduceLInfNormKeepDimModule())
+def ReduceLInfNormKeepDimModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4))
+
+
+# ==============================================================================
+
+
+# ord = -inf reducing over all dims (dim=None) of a multidim input.
+class ReduceLNegInfNormNoneDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=None, ord=float("-inf"))
+
+
+@register_test_case(module_factory=lambda: ReduceLNegInfNormNoneDimModule())
+def ReduceLNegInfNormNoneDimModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4))
+
+
+# ==============================================================================
+
+
+# ord = 0 over multiple dims, with keepdim, of a multidim input.
+class ReduceL0NormMultiDimKeepDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.linalg.vector_norm(a, dim=[0, 1], keepdim=True, ord=0)
+
+
+@register_test_case(module_factory=lambda: ReduceL0NormMultiDimKeepDimModule())
+def ReduceL0NormMultiDimKeepDimModule_basic(module, tu: TestUtils):
+    # Include exact zeros so the nonzero-count semantics are exercised.
+    module.forward(torch.tensor([[0.0, 1.5, 0.0], [-2.0, 0.0, 3.0]]))
+
+
+# ==============================================================================
+
+
 class ReduceL3NormAllDimsModule(torch.nn.Module):
     def __init__(self):
         super().__init__()