)]}'
{
  "commit": "bb3a0d9e90f4a3608d1739761faebd500783d5f5",
  "tree": "17febcf55d1d142d77bdcc8966cd79e79e3d3525",
  "parents": [
    "0f4ceac7a6e3b3dd17552014820575923be070a6"
  ],
  "author": {
    "name": "Yuwei Sun",
    "email": "49397791+YuWei-CH@users.noreply.github.com",
    "time": "Fri May 29 13:22:05 2026 -0500"
  },
  "committer": {
    "name": "GitHub",
    "email": "noreply@github.com",
    "time": "Fri May 29 20:22:05 2026 +0200"
  },
  "message": "[Codegen] Split StableHLO/JAX select-style argmax reductions (#24445)\n\nProblem:\nStableHLO/JAX lowers argmax for lm_head vocab reductions as a pair of\nselects with explicit NaN handling and lower-index tie breaking. The\nexisting split-reduction matcher only recognized the canonical\nmaximumf/cmpf/select form, so large vocab reductions stayed as one wide\nreduction.\n\nFix:\nTeach the argmax matcher to distinguish canonical and StableHLO/JAX\nselect-style argmax reductions, and preserve the corresponding combiner\nwhen splitting. Keep the select-style matcher strict by requiring the\n`gt || isnan` value condition, the `eq \u0026\u0026 lower_index` tie condition,\nand the expected value/index select operands.\n\nExample:\nA Qwen `lm_head` argmax over vocab size 248320 can now split into a\ntwo-stage reduction over 1940x128 tiles while preserving NaN and\nlower-index tie semantics. Negative tests ensure similar-looking select\nreductions with different NaN/tie/index semantics do not split.\n\n---------\n\nSigned-off-by: Yuwei Sun \u003cyuweis2@illinois.edu\u003e\nSigned-off-by: Hao Ren \u003crhao8608@gmail.com\u003e\nCo-authored-by: Hao Ren \u003crhao8608@gmail.com\u003e",
  "tree_diff": [
    {
      "type": "modify",
      "old_id": "adab4243aa2f5b11e3c7536ccda5ae7976ef821c",
      "old_mode": 33188,
      "old_path": "compiler/plugins/target/ROCM/test/config_ukernel_argmax_gfx942.mlir",
      "new_id": "e7cd19e6e4cd58527ae8e1256a2e549207043b95",
      "new_mode": 33188,
      "new_path": "compiler/plugins/target/ROCM/test/config_ukernel_argmax_gfx942.mlir"
    },
    {
      "type": "modify",
      "old_id": "ab7c1a3b5a8f2a58827de7b02b22d5a5df71a5c1",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp",
      "new_id": "5065a4e8a93d807514d2701e2789c11c8e5550cc",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/SplitReduction.cpp"
    },
    {
      "type": "modify",
      "old_id": "932358afa577c2949ca1c4e889051f89a25fba82",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp",
      "new_id": "7d04491b8fae3be7d97962bf33991bc476421207",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp"
    },
    {
      "type": "modify",
      "old_id": "2f90ed052c731b363a19f1166fe1554498a0c48a",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h",
      "new_id": "c3ca59892c0320676013582aec28156ef55ea2af",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
    },
    {
      "type": "modify",
      "old_id": "092ec2ab178847ebb4ad9b82b71d85971831ef51",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir",
      "new_id": "611b4623f2ba8bd4adb1e57a488803d8e9be4b81",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/DispatchCreation/test/split_reduction.mlir"
    }
  ]
}