define_extern

[stfeiseu]

Recently, I tried using the defin_extern method defined by Halide with AOT mode. The description code of the Halide frontend is as follows:

class ExternFuncInput : public Generator<ExternFuncInput> {
    Input<Buffer<int>> input{"input", 2};
    Input<int> addend{"addend"};
    Output<Buffer<int>> output{"output", 2};
    Func input_ub{"input_ub"}; 
    Func work{"work"}; 
    Func extern_func{"extern_func"};
    Func vmuls{"vmuls"};
    Var x{"x"}, y{"y"};

public:
    void generate() {
        input_ub(x, y) = input(x, y);
        work(x, y) = input_ub(x, y) * 2;
        std::vector<ExternFuncArgument> params = {addend};
        std::vector<Type> types = {Int(32)};
        std::vector<Var> args = {x, y};
        extern_func.define_extern("extern_stage", params, types, args, NameMangling::Default, DeviceAPI::CUDA);
        vmuls(x, y) = extern_func(x, y) * 3;
        output(x, y) = vmuls(x, y) + work(x, y);
    }

    void schedule() {
        Var xo, yo, xi, yi;
        // output.tile(x, y, xo, yo, xi, yi, 128, 128);
        output.split(x, xo, xi, 128);
        output.split(y, yo, yi, 128);
        output.reorder(xi, yi, xo, yo);
        input_ub.compute_at(output, xo);   
        work.compute_at(output, xo).vectorize(x, 32);
        extern_func.compute_root();
        vmuls.compute_at(output, xo).vectorize(x, 32);
        output.compute_root().vectorize(xi, 32);
    }
};

After Unrolling Pass, the IR are as follows:

Lowering after simplifying correlated differences:
assert((call)(uint64)reinterpret((Variable)(halide_buffer_t *)output.buffer) != (uint64)0, (call)halide_error_buffer_argument_is_null("output"))
assert((call)(uint64)reinterpret((Variable)(halide_buffer_t *)input.buffer) != (uint64)0, (call)halide_error_buffer_argument_is_null("input"))
let input = (call)(void *)_halide_buffer_get_host((Variable)(halide_buffer_t *)input.buffer)
let input.type = (call)(uint32)_halide_buffer_get_type((Variable)(halide_buffer_t *)input.buffer)
let input.dimensions = (call)_halide_buffer_get_dimensions((Variable)(halide_buffer_t *)input.buffer)
let input.min.0 = (call)_halide_buffer_get_min((Variable)(halide_buffer_t *)input.buffer, 0)
let input.extent.0 = (call)_halide_buffer_get_extent((Variable)(halide_buffer_t *)input.buffer, 0)
let input.stride.0 = (call)_halide_buffer_get_stride((Variable)(halide_buffer_t *)input.buffer, 0)
let input.min.1 = (call)_halide_buffer_get_min((Variable)(halide_buffer_t *)input.buffer, 1)
let input.extent.1 = (call)_halide_buffer_get_extent((Variable)(halide_buffer_t *)input.buffer, 1)
let input.stride.1 = (call)_halide_buffer_get_stride((Variable)(halide_buffer_t *)input.buffer, 1)
let output = (call)(void *)_halide_buffer_get_host((Variable)(halide_buffer_t *)output.buffer)
let output.type = (call)(uint32)_halide_buffer_get_type((Variable)(halide_buffer_t *)output.buffer)
let output.dimensions = (call)_halide_buffer_get_dimensions((Variable)(halide_buffer_t *)output.buffer)
let output.min.0 = (call)_halide_buffer_get_min((Variable)(halide_buffer_t *)output.buffer, 0)
let output.extent.0 = (call)_halide_buffer_get_extent((Variable)(halide_buffer_t *)output.buffer, 0)
let output.stride.0 = (call)_halide_buffer_get_stride((Variable)(halide_buffer_t *)output.buffer, 0)
let output.min.1 = (call)_halide_buffer_get_min((Variable)(halide_buffer_t *)output.buffer, 1)
let output.extent.1 = (call)_halide_buffer_get_extent((Variable)(halide_buffer_t *)output.buffer, 1)
let output.stride.1 = (call)_halide_buffer_get_stride((Variable)(halide_buffer_t *)output.buffer, 1)
let extern_func.o0.bounds_query = let t3 = (call)(halide_dimension_t *)make_struct((Variable)output.min.0, (Variable)output.extent.0, 0, 0, (Variable)output.min.1, (Variable)output.extent.1, 0, 0) in (call)(halide_buffer_t *)_halide_buffer_init((call)(halide_buffer_t *)alloca((call)size_of_halide_buffer_t()), (Variable)t3, (call)(void *)reinterpret((uint64)0), (uint64)0, (call)(halide_device_interface_t *)reinterpret((uint64)0), 0, 32, 2, (Variable)t3, (uint64)0)
let t4 = (call)extern_stage((Variable)addend, (Variable)extern_func.o0.bounds_query)
assert((Variable)t4 == 0, (call)halide_error_bounds_inference_call_failed("extern_stage", (Variable)t4))
let extern_func.s0.y.min_1 = (call)_halide_buffer_get_min((Variable)extern_func.o0.bounds_query, 1)
let extern_func.s0.y.max_1 = (call)_halide_buffer_get_max((Variable)extern_func.o0.bounds_query, 1)
let extern_func.s0.x.min_1 = (call)_halide_buffer_get_min((Variable)extern_func.o0.bounds_query, 0)
let extern_func.s0.x.max_1 = (call)_halide_buffer_get_max((Variable)extern_func.o0.bounds_query, 0)
let extern_func.o0.bounds_query_1 = let t1 = (call)(halide_dimension_t *)make_struct((Variable)extern_func.s0.x.min_1, (add)(sub)((Variable)extern_func.s0.x.max_1 - (Variable)extern_func.s0.x.min_1) + 1, 0, 0, (Variable)extern_func.s0.y.min_1, (add)(sub)((Variable)extern_func.s0.y.max_1 - (Variable)extern_func.s0.y.min_1) + 1, 0, 0) in (call)(halide_buffer_t *)_halide_buffer_init((call)(halide_buffer_t *)alloca((call)size_of_halide_buffer_t()), (Variable)t1, (call)(void *)reinterpret((uint64)0), (uint64)0, (call)(halide_device_interface_t *)reinterpret((uint64)0), 0, 32, 2, (Variable)t1, (uint64)0)
let t2 = (call)extern_stage((Variable)addend, (Variable)extern_func.o0.bounds_query_1)
assert((Variable)t2 == 0, (call)halide_error_bounds_inference_call_failed("extern_stage", (Variable)t2))
let input.extent.0.required.s = (sub)min((mul)(div)((add)((Variable)output.extent.0 + -1)/128)*128, (add)(Variable)output.extent.0 + -128) - min((Variable)output.extent.0, 128)
let input.extent.1.required.s = (sub)min((mul)(div)((add)((Variable)output.extent.1 + -1)/128)*128, (add)(Variable)output.extent.1 + -128) - min((Variable)output.extent.1, 128)
let output.extent.0.required.s = (Variable)input.extent.0.required.s
let output.extent.1.required.s = (Variable)input.extent.1.required.s
...
...
allocate extern_func[...]
(ProducerConsumer)produce extern_func {
  let halide_device_malloc_result = (call)halide_device_malloc((Variable)extern_func.buffer, (call)(halide_device_interface_t *)halide_cuda_device_interface())
  assert((Variable)halide_device_malloc_result == 0, (Variable)halide_device_malloc_result)
  let t0 = (call)extern_stage((Variable)addend, (Variable)extern_func.buffer)
  assert((Variable)t0 == 0, (call)halide_error_extern_stage_failed("extern_stage", (Variable)t0))
  (evaluate)(call)_halide_buffer_set_device_dirty((Variable)extern_func.buffer, (uint1)1)
 }

My question is :
Why is extern_stage called multiple times before extern_stage is generated? Can anyone help me? Thank you very much.

[abadams]

Halide needs to know how much input your extern stage needs for a given tile of output in order to be able to compute stuff before the extern stage in your pipeline. It calls the stage with buffers with null host pointers to make these queries. The idea is that the stage is supposed to recognize this, and fill in the buffers appropriately. See here for an example:

https://github.com/halide/Halide/blob/master/test/correctness/extern_stage.cpp

So the extra calls are what we call "bounds queries".

[stfeiseu]

Andrew, I'm very grateful for your answer. Based on your answer and the example of the extern function that comes with the Halide source code. My understanding of the extern function is as follows:

1. In Halide IR, the extern function is called at least once before the extern stage in the pipeline, regardless of whether it contains input parameters.
2. If the input parameters of the extern function include buffer and the input buffer is_bounds_query, the size relationship between the specified input buffer and output buffer must be displayed during the implementation of the extern function so that Halide can query the size of the input buffer.

If there is any mistake in my understanding of the two above, please correct it.
Thank you very much.