diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir index 1a213af6ae2e..640da7802f1a 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir +++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir @@ -139,14 +139,14 @@ util.func public @tensorSplat(%value: i8, %dim0: index) -> tensor { util.global private @device : !hal.device // CHECK-LABEL: @tensorBarrierDispatch -// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[DIM0:.+]]: index, %[[DIM1:.+]]: index) +// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index, %[[DIM0:.+]]: index) util.func public @tensorBarrierDispatch(%input: tensor, %dim0: index) -> tensor { - // CHECK: %[[BARRIER:.+]] = stream.async.barrier on(#hal.device.affinity<@device>) %[[INPUT]] : !stream.resource<*>{%[[DIM0]]} + // CHECK: %[[BARRIER:.+]] = stream.async.barrier on(#hal.device.affinity<@device>) %[[INPUT]] : !stream.resource<*>{%[[INPUT_SIZE]]} %barrier = flow.tensor.barrier %input : tensor{%dim0} on #hal.device.affinity<@device> - // CHECK: %[[SIZE:.+]] = stream.tensor.sizeof on(#hal.device.affinity<@device>) tensor{%arg2} : index - // CHECK: %[[RESULT:.+]] = stream.tensor.dispatch on(#hal.device.affinity<@device>) @ex::@entry(%[[BARRIER]]) + // CHECK: %[[RESULT_SIZE:.+]] = stream.tensor.sizeof tensor{%[[DIM0]]} : index + // CHECK: %[[RESULT:.+]] = stream.tensor.dispatch @ex::@entry(%[[BARRIER]]) %0 = flow.dispatch @ex::@entry(%barrier) : (tensor{%dim0}) -> tensor{%dim0} - // CHECK: util.return %[[RESULT]], %[[SIZE]] + // CHECK: util.return %[[RESULT]], %[[RESULT_SIZE]] util.return %0 : tensor }