From 50560b919e9b56c04feb2b72ddfc244a72738661 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 27 May 2026 11:43:05 +0900 Subject: [PATCH] [Frontend] Drop stale single-reduction-axis guard breaking test_diffusion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guard at mlir_codegen_backend.py:909-910 if len(reductions.loops) > 1: raise NotImplementedError("Not support multiple reduction axis..") was a leftover from when the codegen was deliberately limited to a single reduction axis, but multi-axis support was restored months ago and the guard was never removed. Timeline: - 8eacf27 ("Optimize reduce/elementwise code", 2025): replaced `for reduction in reductions.loops` with `reductions.loops[0]`, added this guard (no `raise` — silently dead). - c61f67d ("Support multi reduction dim + Add Diffusion model test", 2025-08-14): restored the `for reduction_loop in reductions.loops` iteration explicitly to support multi-axis reduction. test_diffusion was added in this same commit and passed. The guard was left in place — harmless because still `raise`-less. - 5045837 ("Fix four codegen correctness issues surfaced by review", 2026-05-26, issue #237): added the missing `raise`, framing it as a "silently fell through to incorrect codegen" fix. In fact the codegen below the guard (which iterates all reduction loops) was correct — c61f67d had made it so — and adding `raise` broke the very test (test_diffusion) that c61f67d added to validate multi-axis support. Right fix is to delete the guard, not add `raise` to it. The other three items in #237 are unrelated and stand. Verified: file py_compiles. CI test_diffusion is the empirical proof that the post-guard codegen handles multi-axis correctly within the allclose tolerance the test uses. Co-Authored-By: Claude Opus 4.7 (1M context) --- PyTorchSimFrontend/mlir/mlir_codegen_backend.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 9c20311c..b163ad1a 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -906,9 +906,6 @@ def codegen_loops(self): if (self.reduction_depth==0): loops = LoopNest([LoopLevel("dummy", 1)]) - if len(reductions.loops) > 1: - raise NotImplementedError("Not support multiple reduction axis..") - code.splice(self.const_buffer) code.splice(self.alloc_buffer) code.splice(self.spad_buffer)