Halide Python API — поддержка формата чередующихся изображений ⇐ Python

[Anonymous]

Я написал простой генератор Halide на Python и скомпилировал его с помощью webAssembly. Я хочу, чтобы в качестве входных и выходных данных использовались чередующиеся данные изображения RGB. Вот как это выглядит:

Код: Выделить всё

import halide as hl

x = hl.Var('x')
y = hl.Var('y')
c = hl.Var('c')

@hl.generator(name="sample_generator")
class SampleGenerator:
input = hl.InputBuffer(hl.UInt(8), 3)
output = hl.OutputBuffer(hl.UInt(8), 3)

def generate(self):
self.output[x, y, c] = self.input[x, y, c]

# Set buffer strides for input
self.input.dim(0).set_stride(3)
self.input.dim(2).set_stride(1)

# Set buffer strides for output
#   Test 1. Does not work since OutputBuffer does not have dim()
# self.output.dim(0).set_stride(3)
# self.output.dim(2).set_stride(1)
#   Test 2. Set stride on output_buffer(). Runs fine, but still complain at runtime on contraint violated on output.stride.0 (3) == 1
self.output.output_buffer().dim(0).set_stride(3)
self.output.output_buffer().dim(2).set_stride(1)

В вызове JS я создаю halide_buffer_t с тремя измерениями (например, (256, 256, 3) с шагом (3, 256*3, 1). Код для создания этого halide_buffer_t доступен на C++ для веб-сборки:

Код: Выделить всё

halide_buffer_t* make_buffer(
void *data,
halide_type_t type,
int dimensions,
halide_dimension_t *dims
) {
halide_buffer_t* buf = new halide_buffer_t{};
buf->device = 0;
buf->host = static_cast(data);
buf->flags = 0;
buf->type = type;
buf->dimensions = dimensions;
buf->dim = dims;
buf->padding = nullptr;
return buf;
}

template  halide_buffer_t* make_buffer(void* data, int dimensions, halide_dimension_t *dims) {
// Just create a halide_type_t and forward dims
// ...
return make_buffer(data, type, dimensions, dims);
}

template 
halide_buffer_t* make_buffer_3d_interleaved(
void *data,
int width,
int height,
int channels
) {
auto *dims = new halide_dimension_t[3];

dims[0].min = 0;
dims[0].extent = width;
dims[0].stride = channels;
dims[0].flags = 0;

dims[1].min = 0;
dims[1].extent = height;
dims[1].stride = width * channels;
dims[1].flags = 0;

dims[2].min = 0;
dims[2].extent = channels;
dims[2].stride = 1;
dims[2].flags = 0;

return make_buffer(data, 3, dims);
}

Как видите, я попробовал несколько способов задать ограничения шага для Halide в генераторе, но безуспешно.
В C++ это мой обычный способ установки таких ограничений. Я также заметил, что если я запускаю генератор из Python, используя массивы numpy в качестве входных данных, все идет нормально, даже если я не устанавливаю никаких ограничений на шаги данных, поэтому я считаю, что среда выполнения Python Halide сделает свое дело.
Как правильно использовать чередующиеся форматы в этом случае?

Спасибо

Подробнее здесь: https://stackoverflow.com/questions/798 ... at-support