hal.fence.join: OUT_OF_RANGE; count 90 of iree_hal_fence > 32

[silvasean]

What happened?

Using ir19.flagfile and ir19.no_sharding.mlir, run the following commands:

iree-compile --iree-hal-target-backends=cuda --iree-input-type=mhlo --iree-hal-cuda-llvm-target-arch=sm_80 ir19.no_sharding.mlir --iree-codegen-llvmgpu-use-transform-dialect= --iree-codegen-llvmgpu-enable-transform-dialect-jit=false -o ir19.vmfb
iree-benchmark-module --function=main --module=ir19.vmfb --device=cuda --flagfile=ir19.flagfile --benchmark_repetitions=5

(The --iree-codegen-llvmgpu-use-transform-dialect= --iree-codegen-llvmgpu-enable-transform-dialect-jit=false is to work around #13419)

I see the following error:

iree/runtime/src/iree/modules/hal/module.c:1119: OUT_OF_RANGE; count 90 of iree_hal_fence > 32; while invoking native function hal.fence.join; while calling import; 
.....

(full log here)

Steps to reproduce your issue

See above

What component(s) does this issue relate to?

Runtime

Version information

iree.git @ 4d6c2b8 (I have a couple trivial local modifications, but should be unrelated to this bug)

Additional context

No response

[silvasean]

Is the solution here as simple as just bumping up the number from 32 to 128 or something like that?

[benvanik]

We could slightly up it as a whackamole smack - it's guarding a stack allocation though and we'd want to keep it small and special case anything larger with a heap allocation. I think the real issue is that you have 90 non-folded fences - those are (relatively) expensive and indicate some earlier stage of the pipeline being really inefficient. --compile-to=stream and post the IR and we can see if it's obvious.

[allieculp]

@pavanimajety Curious if there is any update on this?

[silvasean]

@benvanik here is the --compile-to=stream output you requested: https://gist.github.com/silvasean/1038fec28172cf61872cdd6e000523ad

[silvasean]

For context, this is a LLM training workload. The IR structure is the forward pass, backward pass, and then optimizer updates. The forward and backward pass each have their own stablehlo.while op. In the current IR, this actually only iterates once (note sure if we fold it away, I see a ton of control flow in the --compile-to=stream IR, more than I would expect for this). In the user workload this corresponds to microbatching (see the image here).

[allieculp]

@pjannaty @pavanimajety Is this being worked on by your team?

[pavanimajety]

@allieculp I am taking a look, was out of office last few days.

[pjannaty]

Thanks, Allie and Pavani. Allie, this looks like a high-priority bug and we are gradually ramping up and don't mean to block if someone else also wants to take a look. Happy to collaborate on this as well.

[benvanik]

issue is FoldBlockArgumentsPattern not properly folding the duplicate block arguments. Will see if I can get that fixed.