WhileThunk may block in async mode

[wjr-z]

xla/xla/backends/cpu/runtime/while_thunk.cc

Lines 92 to 101 in eb9c4f6

	for (int64_t loop_counter = 0; loop_counter < trip_count; ++loop_counter) {
	auto body_event = body_executor_.Execute(params);
	// If loop iteration has not completed yet, switch to async execution mode
	// using `body_event` as a dependency and continue the loop iteration
	// starting from `loop_counter + 1`.
	if (ABSL_PREDICT_FALSE(!body_event.IsAvailable())) {
	return ExecuteAsyncForLoop(params, std::move(body_event),
	loop_counter + 1, trip_count);
	}

Sometimes even if the body is fast enough to execute synchronously, it may block for a long time when the number of iterations is large enough.

Possible solutions:

1. Use cost_model of body function(If possible) to decide whether to execute asynchronously.
2. Use timeout timer to execute asynchronously after exceeding the threshold.

[wjr-z]

Use AndThen here may still block in that cast

[wjr-z]

Slow down in my case if i set ExecuteMayBlock to true when schdulable CPU number is 128, but faster when i set CPU number to 2~4.
In my example, a while is used to process a 2000 * 4000 matrix to do some operations.

[penumallas1261]

Hi @wjr-z ,

Thank you for creating this issue. To help us investigate, please share the steps you performed and simple standalone code snippet to reproduce the problem.

[wjr-z]

嗨，

感谢您创建此问题。为了帮助我们进行调查，请分享您执行的步骤以及用于重现问题的简单独立代码片段。

@penumallas1261

My use case is quite complex, so I am trying to simplify it.
During this process, I discovered another issue

Deps of ExecutionGraph in my case is abnormal.
Here's my test cast:

#include <iterator>
#include <chrono>

#include "xla/hlo/builder/xla_builder.h"
#include "xla/client/client.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/pjrt/cpu/cpu_client.h"
#include "xla/pjrt/cpu/cpu_device.h"

const int64_t n = 32000;
const int64_t m = 4000;

xla::XlaOp BuildWhile(xla::XlaBuilder& builder, xla::XlaOp input) {
    auto input_shape = xla::ShapeUtil::MakeShape(xla::F64, {n, m});

    auto init_row_idx = xla::ConstantR0<int64_t>(&builder, 0);
    auto init_sum = xla::Broadcast(xla::ConstantR0<double>(&builder, 0.0), {1, m});

    auto init_tuple_with_input = xla::Tuple(&builder, {init_row_idx, init_sum, input});
    auto tuple_shape = xla::ShapeUtil::MakeTupleShape(
        {xla::ShapeUtil::MakeShape(xla::S64, {}),
         xla::ShapeUtil::MakeShape(xla::F64, {1, m}), input_shape});

    xla::XlaComputation condition_with_input;
    {
        xla::XlaBuilder cond_builder("condition_with_input");

        auto param = xla::Parameter(&cond_builder, 0, tuple_shape, "state");
        auto row_idx = xla::GetTupleElement(param, 0);
        auto n_const = xla::ConstantR0<int64_t>(&cond_builder, n);
        xla::Lt(row_idx, n_const);
        condition_with_input = cond_builder.Build().value();
    }

    xla::XlaComputation body_with_input;
    {
        xla::XlaBuilder body_builder("body_with_input");
        auto param = xla::Parameter(&body_builder, 0, tuple_shape, "state");

        auto row_idx = xla::GetTupleElement(param, 0);
        auto sum_vec = xla::GetTupleElement(param, 1);
        auto input_param = xla::GetTupleElement(param, 2);

        auto current_row = xla::DynamicSlice(
            input_param, {row_idx, xla::ConstantR0<int64_t>(&body_builder, 0)}, {1, m});

        auto new_sum = xla::Add(sum_vec, current_row);

        auto next_idx = xla::Add(row_idx, xla::ConstantR0<int64_t>(&body_builder, 1));

        xla::Tuple(&body_builder, {next_idx, new_sum, input_param});
        body_with_input = body_builder.Build().value();
    }

    auto result_tuple =
        xla::While(condition_with_input, body_with_input, init_tuple_with_input);

    auto result = xla::GetTupleElement(result_tuple, 1);
    return result;
}

xla::XlaComputation RowSumWithWhile(size_t K) {
    xla::XlaBuilder builder("row_sum_while");

    auto input_shape = xla::ShapeUtil::MakeShape(xla::F64, {n, m});

    std::vector<xla::XlaOp> outputs;
    for (size_t i = 0; i < K; ++i) {
        auto input = xla::Parameter(&builder, i, input_shape, "input");
        outputs.push_back(BuildWhile(builder, input));
    }

    xla::Tuple(&builder, outputs);
    return builder.Build().value();
}

#include <fmt/format.h>

int main(int argc, char** argv) {
    size_t K = 1;
    if (argc > 1) {
        K = std::stoul(argv[1]);
    }

    std::cout << "K = " << K << std::endl;

    auto client = xla::GetPjRtCpuClient({}).value();

    auto computation = RowSumWithWhile(K);

    xla::CompileOptions options;
    auto& build_options = options.executable_build_options;
    build_options.set_optimization_level(xla::ExecutionOptions::EFFORT_O3);

    auto pjrt_executable =
        client->CompileAndLoad(computation, std::move(options)).value();

    std::vector<double> input_data(n * m, 1.0);
    std::vector<std::unique_ptr<xla::PjRtBuffer>> input_buffers;
    std::vector<xla::PjRtBuffer*> buffers;
    for (size_t i = 0; i < K; ++i) {
        input_buffers.push_back(
            client
                ->BufferFromHostBuffer(
                    input_data.data(), xla::F64, {n, m}, std::nullopt,
                    xla::PjRtClient::HostBufferSemantics::kImmutableZeroCopy, []() {},
                    client->devices().at(0)->default_memory_space().value(), nullptr)
                .value());
        buffers.push_back(input_buffers.back().get());
    }

    auto start = std::chrono::high_resolution_clock::now();

    xla::ExecuteOptions exec_options;
    auto result_or =
        pjrt_executable->Execute({std::move(buffers)}, std::move(exec_options));
    if (!result_or.ok()) {
        return 1;
    }
    auto result = std::move(result_or.value());
    for (auto& buffer : result.at(0)) {
        buffer->GetReadyFuture().Await();
    }

    auto end = std::chrono::high_resolution_clock::now();
    std::cout
        << std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count()
        << std::endl;
    return 0;
}

BuildWhile is a simple While Loop to add every row from a input NxM matrix to a output 1xM vector.
The k-th BuildWhile accept k-th parmeter as input.

Run program with arg 2.

By printing ThunkExecutor ToString, I found dependency relationship between two While Thunk.
I printed BufferUse and found overlap size=1 and size=8, this should be easy to avoid?

As for the blocking of While Thunk itself, I'm currently trying to build tests

[wjr-z]

Log of ThunkExecutor

I0000 00:00:1761290278.595411 2406972 thunk_executor.cc:169] ThunkExecutor execution graph:
ThunkExecutor: #thunks=7 #source_nodes=3 #sink_nodes=1
 thunk #00000: op_name=copy, kind=copy, dependencies=[], source=true, sink=false, priority=2
 thunk #00001: op_name=copy.6, kind=copy, dependencies=[], source=true, sink=false, priority=1
 thunk #00002: op_name=wrapped_broadcast, kind=kernel, dependencies=[], source=true, sink=false, priority=4
 thunk #00003: op_name=copy.1, kind=copy, dependencies=[wrapped_broadcast], source=false, sink=false, priority=2
 thunk #00004: op_name=copy.7, kind=copy, dependencies=[wrapped_broadcast], source=false, sink=false, priority=2
 thunk #00005: op_name=while.2, kind=while, dependencies=[copy, copy.1, copy.7], source=false, sink=false, priority=1
 thunk #00006: op_name=while.3, kind=while, dependencies=[copy.6, while.2], source=false, sink=true, priority=0

Log of BufferUses

I0000 00:00:1761290278.595304 2406972 execution_graph.cc:133] Operation 0 Operation {name: op: copy (kind: copy)}
Buffer uses:
{slice: {index:7, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:32192, size:8}, access: W, content_validity: O}

I0000 00:00:1761290278.595329 2406972 execution_graph.cc:133] Operation 1 Operation {name: op: copy.6 (kind: copy)}
Buffer uses:
{slice: {index:7, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:32128, size:8}, access: W, content_validity: O}

I0000 00:00:1761290278.595335 2406972 execution_graph.cc:133] Operation 2 Operation {name: op: wrapped_broadcast (kind: kernel)}
Buffer uses:
{slice: {index:8, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:32000}, access: W, content_validity: O}

I0000 00:00:1761290278.595340 2406972 execution_graph.cc:133] Operation 3 Operation {name: op: copy.1 (kind: copy)}
Buffer uses:
{slice: {index:11, offset:128, size:32000}, access: R, content_validity: IO}
{slice: {index:2, offset:0, size:32000}, access: W, content_validity: O}

I0000 00:00:1761290278.595344 2406972 execution_graph.cc:133] Operation 4 Operation {name: op: copy.7 (kind: copy)}
Buffer uses:
{slice: {index:11, offset:128, size:32000}, access: R, content_validity: IO}
{slice: {index:3, offset:0, size:32000}, access: W, content_validity: O}

I0000 00:00:1761290278.595353 2406972 execution_graph.cc:133] Operation 5 Operation {name: op: while.2 (kind: while)}
Buffer uses:
{slice: {index:11, offset:128, size:1}, access: W, content_validity: O}
{slice: {index:11, offset:32192, size:8}, access: R, content_validity: IO}
{slice: {index:5, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:1}, access: W, content_validity: O}
{slice: {index:11, offset:32192, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:8}, access: W, content_validity: O}
{slice: {index:11, offset:128, size:8}, access: R, content_validity: IO}
{slice: {index:4, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:32192, size:8}, access: W, content_validity: O}
{slice: {index:2, offset:0, size:32000}, access: R, content_validity: IO}
{slice: {index:0, offset:0, size:1024000000}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:8}, access: R, content_validity: IO}
{slice: {index:2, offset:0, size:32000}, access: W, content_validity: O}

I0000 00:00:1761290278.595370 2406972 execution_graph.cc:133] Operation 6 Operation {name: op: while.3 (kind: while)}
Buffer uses:
{slice: {index:11, offset:128, size:1}, access: W, content_validity: O}
{slice: {index:11, offset:32128, size:8}, access: R, content_validity: IO}
{slice: {index:9, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:1}, access: W, content_validity: O}
{slice: {index:11, offset:32128, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:8}, access: W, content_validity: O}
{slice: {index:11, offset:128, size:8}, access: R, content_validity: IO}
{slice: {index:10, offset:0, size:8}, access: R, content_validity: IO}
{slice: {index:11, offset:32128, size:8}, access: W, content_validity: O}
{slice: {index:3, offset:0, size:32000}, access: R, content_validity: IO}
{slice: {index:1, offset:0, size:1024000000}, access: R, content_validity: IO}
{slice: {index:11, offset:128, size:8}, access: R, content_validity: IO}
{slice: {index:3, offset:0, size:32000}, access: W, content_validity: O}

[wjr-z]

xla/xla/hlo/utils/hlo_live_range.cc

Lines 255 to 257 in 6048a86

	} else {
	live_range.end = std::max(live_range.end, GetLastUsageTime(*value));
	}

If I write this below:

} else {
    live_range.end = schedule_end_time();
}

Then there's no deps between two WhileOp.
This place, extend live range if value is used.
Return value of WhileOp's condition computation is used by WhileOp, but its GetUses() not contains WhileOp.

xla/xla/hlo/utils/hlo_live_range.cc

Lines 169 to 171 in 6048a86

	LogicalTime end_time = -1;
	for (const HloUse& use : value.GetUses()) {
	const HloInstruction* used = use.instruction;

So live range of WhileOp's condition result value maybe end early. Then maybe will conflict with others.

Additionally, the performance testing here may not be very accurate as I did not use 64 byte aligned vector, resulting in a significant loss of performance due to copying.
After using my own memory allocator here and making While parallelizable, I achieved a good acceleration ratio