stablehlo.concatenate: error: operand #4 does not dominate this use

[MichelBartels]

What happened?

When compiling a random number generator to StableHLO and then executing it using IREE, IREE is complaining about an operand #4 to "stablehlo.concatenate" not dominating its uses. I managed to reduce it to a simple program where it is obvious that the error message is wrong.

I am also not sure what operand #4 is referring to as there is no fourth operand.

Steps to reproduce your issue

bug.mlir (simplest case I could reduce it to):

func.func @main(%0 : tensor<1xui32>, %1: tensor<1xf32>) -> (tensor<2xf32>, tensor<1xui32>) {
%3 = "stablehlo.add"(%0, %0) : (tensor<1xui32>, tensor<1xui32>) -> (tensor<1xui32>)
%4 = "stablehlo.convert"(%3) : (tensor<1xui32>) -> (tensor<1xf32>)
%5 = "stablehlo.negate"(%4) : (tensor<1xf32>) -> (tensor<1xf32>)
%6 = "stablehlo.add"(%1, %1) : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
%7 = "stablehlo.concatenate"(%4, %6) {
dimension = 0
} : (tensor<1xf32>, tensor<1xf32>) -> (tensor<2xf32>)
%8 = "stablehlo.add"(%0, %3) : (tensor<1xui32>, tensor<1xui32>) -> (tensor<1xui32>)
"stablehlo.return"(%7, %8) : (tensor<2xf32>, tensor<1xui32>) -> ()
}

Running iree-compile bug.mlir results in the following error message:

bug.mlir:6:6: error: operand #4 does not dominate this use
%7 = "stablehlo.concatenate"(%4, %6) {
     ^
bug.mlir:1:1: note: called from
func.func @main(%0 : tensor<1xui32>, %1: tensor<1xf32>) -> (tensor<2xf32>, tensor<1xui32>) {
^
bug.mlir:6:6: note: see current operation: %12 = "stream.async.update"(%11, %0, %3, %1, %13#2, %1) : (!stream.resource<*>, index, index, index, !stream.resource<*>, index) -> !stream.resource<*>
%7 = "stablehlo.concatenate"(%4, %6) {
     ^
bug.mlir:5:6: note: operand defined here (op in the same block)
%6 = "stablehlo.add"(%1, %1) : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
     ^

What component(s) does this issue relate to?

No response

Version information

Version 3.0.0 of iree-compiler-base from pypi

Additional context

No response

[MichelBartels]

Reordering the instructions gets rid of the error message, but I don't see why this should make a difference.

func.func @main(%0 : tensor<1xui32>, %1: tensor<1xf32>) -> (tensor<2xf32>, tensor<1xui32>) {
%6 = "stablehlo.add"(%1, %1) : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
%3 = "stablehlo.add"(%0, %0) : (tensor<1xui32>, tensor<1xui32>) -> (tensor<1xui32>)
%4 = "stablehlo.convert"(%3) : (tensor<1xui32>) -> (tensor<1xf32>)
%5 = "stablehlo.negate"(%4) : (tensor<1xf32>) -> (tensor<1xf32>)
%7 = "stablehlo.concatenate"(%4, %6) {
dimension = 0
} : (tensor<1xf32>, tensor<1xf32>) -> (tensor<2xf32>)
%8 = "stablehlo.add"(%0, %3) : (tensor<1xui32>, tensor<1xui32>) -> (tensor<1xui32>)
"stablehlo.return"(%7, %8) : (tensor<2xf32>, tensor<1xui32>) -> ()
}

[IanWood1]

Here is a reproducer that can be run with iree-opt --run-reproducer repro.mlir. The problematic pass is iree-stream-emplace-allocations. It seems to be caused by stream.async.update ops having a producer/consumer relationship. Possibly related: #14566

The pass iterates over all of the results to find updates ops to move. In the modified IR, the dispatch op is the last result so all the updates get correctly moved VS the original where only 1 of the 2 are moved. The following seems to be the cause of the difference between the 2 orderings:

iree/compiler/src/iree/compiler/Dialect/Stream/Transforms/EmplaceAllocations.cpp

Lines 83 to 85 in f27feff

	auto updateOp = dyn_cast_or_null<IREE::Stream::AsyncUpdateOp>(userOp);
	if (!updateOp)
	break;

.