`torch.split` compiles with `torch==2.5.1`, but fails with `torch==2.6.0`

[stbaione]

Error

I was working on implementing a custom split argmax op and found that torch.split fails with torch == 2.6.0, but passes with torch==2.5.1.

Artifacts

• split_op_torch_2_5_1.mlir PASSES
• split_op_torch_2_6_0.mlir FAILS

Dependency Versions

iree-turbine==3.4.0rc20250430 (nightly)
iree==d916c3771fb0ef320da2f2f7efc91c1a59c9a962 (built from source)

Failure w/ Torch 2.6.0

With torch==2.6.0, I get the following error when compiling:

Compile Command:

iree-compile split_op_torch_2_6_0.mlir \
    -o split_op_torch_2_6_0.vmfb \
    --iree-hal-target-device=hip \
    --iree-hip-target=gfx942

Error:

<unknown>:0: error: failed to legalize operation 'torch.constant.int'
<unknown>:0: note: see current operation: %0 = "torch.constant.int"() <{value = 8192 : i64}> : () -> !torch.int

When adding --mlir-print-skip-regions,--mlir-print-ir-before-all,--mlir-disable-threading, it seems to be failing on the FinalizingBackendTypeConversion pass:

// -----// IR Dump Before FinalizingBackendTypeConversion (torch-finalizing-backend-type-conversion) //----- //
util.func public @split$async(%arg0: !hal.buffer_view, %arg1: !hal.fence, %arg2: !hal.fence) -> (!hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) attributes {inlining_policy = #util.inline.never, iree.abi.model = "coarse-fences", iree.abi.stub} {...}

<unknown>:0: error: failed to legalize operation 'torch.constant.int'
<unknown>:0: note: see current operation: %0 = "torch.constant.int"() <{value = 8192 : i64}> : () -> !torch.int
// -----// IR Dump Before FinalizingBackendTypeConversion (torch-finalizing-backend-type-conversion) //----- //
util.func public @split(%arg0: !hal.buffer_view) -> (!hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view, !hal.buffer_view) attributes {iree.abi.stub} {...}

This compiles successfully when using torch==2.5.1.