)]}' { "commit": "21c6cd2649e6d060c95af90a715eab3a42e83d95", "tree": "57ffb736f994108c230ee500cf7b96c4a59b90b1", "parents": [ "100a1d393c13d688d2971a3f8b6114bec69fa345" ], "author": { "name": "pstarkcdpr", "email": "paul.stark@cdprojektred.com", "time": "Fri Jul 31 12:18:09 2026 -0700" }, "committer": { "name": "GitHub", "email": "noreply@github.com", "time": "Fri Jul 31 21:18:09 2026 +0200" }, "message": "[Preprocessing] Handle batching dims and dynamic-update-slice scatter ops (#24720)\n\nSee bug report https://github.com/iree-org/iree/issues/24719\n\nJAX-emitted scatters use two forms that IREE\u0027s StableHLO scatter\npreprocessing did not support, causing malformed shape errors during\ninput conversion:\n\n1. Scatter batching dims were only passed through, so the downstream\ncanonicalizers collapsed batching dims together with the real scatter\nloop dim and miscompiled. Added `ScatterBatchingDimsExpander` (benefit\n2). This is ported with just a few style tweaks directly from upstream\nStableHLO\u0027s `ScatterWithBatchingDimsExpander`. I verified that the\nStableHLO pattern wasn\u0027t accessible to be called directly. The only way\nto access it is to run the whole set of compatibility expanders which\nseemed brittle since new expanders can be added over time, so just\ncopied it and its helper functions. I didn\u0027t spend time to fully\nunderstand the code though since it\u0027s an already reviewed pattern used\nto handle this exact case for backwards compatibility in StableHLO.\nBasically, it lowers each batching dim to an explicit iota index column\nso no batching dims remain for the existing patterns. Here\u0027s an example\n```mlir\nmodule @m {\n func.func @s(%operand: tensor\u003c8x33xi32\u003e, %indices: tensor\u003c8x32x1xi32\u003e, %updates: tensor\u003c8x32xi32\u003e) -\u003e tensor\u003c8x33xi32\u003e {\n %0 \u003d \"stablehlo.scatter\"(%operand, %indices, %updates) \u003c{scatter_dimension_numbers \u003d #stablehlo.scatter\u003cinserted_window_dims \u003d [1], input_batching_dims \u003d [0], scatter_indices_batching_dims \u003d [0], scatter_dims_to_operand_dims \u003d [1], index_vector_dim \u003d 2\u003e}\u003e ({\n ^bb0(%arg2: tensor\u003ci32\u003e, %arg3: tensor\u003ci32\u003e):\n %1 \u003d stablehlo.minimum %arg2, %arg3 : tensor\u003ci32\u003e\n stablehlo.return %1 : tensor\u003ci32\u003e\n }) : (tensor\u003c8x33xi32\u003e, tensor\u003c8x32x1xi32\u003e, tensor\u003c8x32xi32\u003e) -\u003e tensor\u003c8x33xi32\u003e\n return %0 : tensor\u003c8x33xi32\u003e\n }\n}\n```\ngets preprocessed into\n```mlir\nmodule @m {\n func.func @s(%arg0: tensor\u003c8x33xi32\u003e, %arg1: tensor\u003c8x32x1xi32\u003e, %arg2: tensor\u003c8x32xi32\u003e) -\u003e tensor\u003c8x33xi32\u003e {\n %0 \u003d stablehlo.iota dim \u003d 0 : tensor\u003c8x32x1xi32\u003e\n %1 \u003d stablehlo.concatenate %0, %arg1, dim \u003d 2 : (tensor\u003c8x32x1xi32\u003e, tensor\u003c8x32x1xi32\u003e) -\u003e tensor\u003c8x32x2xi32\u003e\n %collapsed \u003d tensor.collapse_shape %1 [[0, 1], [2]] : tensor\u003c8x32x2xi32\u003e into tensor\u003c256x2xi32\u003e\n %collapsed_0 \u003d tensor.collapse_shape %arg2 [[0, 1]] : tensor\u003c8x32xi32\u003e into tensor\u003c256xi32\u003e\n %2 \u003d \"stablehlo.scatter\"(%arg0, %collapsed, %collapsed_0) \u003c{indices_are_sorted \u003d false, scatter_dimension_numbers \u003d #stablehlo.scatter\u003cinserted_window_dims \u003d [0, 1], scatter_dims_to_operand_dims \u003d [0, 1], index_vector_dim \u003d 1\u003e, unique_indices \u003d false}\u003e ({\n ^bb0(%arg3: tensor\u003ci32\u003e, %arg4: tensor\u003ci32\u003e):\n %3 \u003d stablehlo.minimum %arg3, %arg4 : tensor\u003ci32\u003e\n stablehlo.return %3 : tensor\u003ci32\u003e\n }) : (tensor\u003c8x33xi32\u003e, tensor\u003c256x2xi32\u003e, tensor\u003c256xi32\u003e) -\u003e tensor\u003c8x33xi32\u003e\n return %2 : tensor\u003c8x33xi32\u003e\n }\n}\n```\n\n2. Single-index, full-rank overwrite scatters are dynamic-update-slice\nsemantics and cannot be represented as `iree_linalg_ext.scatter`. Added\n`ScatterToDynamicUpdateSlice` (benefit 3) to rewrite them into\n`stablehlo.dynamic_update_slice`.\nHere\u0027s an example:\n```mlir\nmodule @m {\n func.func @s(%operand: tensor\u003c8x7x24x32xf32\u003e, %indices: tensor\u003c1xi32\u003e, %updates: tensor\u003c8x4x24x32xf32\u003e) -\u003e tensor\u003c8x7x24x32xf32\u003e {\n %0 \u003d \"stablehlo.scatter\"(%operand, %indices, %updates) \u003c{indices_are_sorted \u003d true, scatter_dimension_numbers \u003d #stablehlo.scatter\u003cupdate_window_dims \u003d [0, 1, 2, 3], scatter_dims_to_operand_dims \u003d [1]\u003e, unique_indices \u003d true}\u003e ({\n ^bb0(%arg2: tensor\u003cf32\u003e, %arg3: tensor\u003cf32\u003e):\n stablehlo.return %arg3 : tensor\u003cf32\u003e\n }) : (tensor\u003c8x7x24x32xf32\u003e, tensor\u003c1xi32\u003e, tensor\u003c8x4x24x32xf32\u003e) -\u003e tensor\u003c8x7x24x32xf32\u003e\n return %0 : tensor\u003c8x7x24x32xf32\u003e\n }\n}\n```\ngets preprocessed into\n```mlir\nmodule @m {\n func.func @s(%arg0: tensor\u003c8x7x24x32xf32\u003e, %arg1: tensor\u003c1xi32\u003e, %arg2: tensor\u003c8x4x24x32xf32\u003e) -\u003e tensor\u003c8x7x24x32xf32\u003e {\n %c \u003d stablehlo.constant dense\u003c0\u003e : tensor\u003ci32\u003e\n %0 \u003d stablehlo.slice %arg1 [0:1] : (tensor\u003c1xi32\u003e) -\u003e tensor\u003c1xi32\u003e\n %1 \u003d stablehlo.reshape %0 : (tensor\u003c1xi32\u003e) -\u003e tensor\u003ci32\u003e\n %2 \u003d stablehlo.dynamic_update_slice %arg0, %arg2, %c, %1, %c, %c : (tensor\u003c8x7x24x32xf32\u003e, tensor\u003c8x4x24x32xf32\u003e, tensor\u003ci32\u003e, tensor\u003ci32\u003e, tensor\u003ci32\u003e, tensor\u003ci32\u003e) -\u003e tensor\u003c8x7x24x32xf32\u003e\n return %2 : tensor\u003c8x7x24x32xf32\u003e\n }\n}\n```\n\nBoth patterns run before the existing scatter canonicalizers, because\nthe first expands out functionality not supported by later patterns, and\nthe second one removes the scatter entirely. Lit tests added for both.\n\n---------\n\nSigned-off-by: Paul Stark \u003cpaul.stark@cdprojektred.com\u003e\nCo-authored-by: Claude Opus 4.8 (1M context) \u003cnoreply@anthropic.com\u003e", "tree_diff": [ { "type": "modify", "old_id": "93f5c4bd7e55d41e584e55aa52ec3df6444bd299", "old_mode": 33188, "old_path": "compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp", "new_id": "e152d9cbfe96fe4d9495b5c1379d3ebcc3d01b7f", "new_mode": 33188, "new_path": "compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp" }, { "type": "modify", "old_id": "97bcedb9d1ecbb245f70ab452fa6d0862967373e", "old_mode": 33188, "old_path": "compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir", "new_id": "3d06187a8c1c5026c90093a2a20b9a606fcff95e", "new_mode": 33188, "new_path": "compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir" } ] }