File size: 1,286 Bytes
9aa90e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | // SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
#pragma once
#include <cstdint>
#include <vector>
namespace tt_dit_planar {
// Per-shard view of an input shard
struct ShardView {
const uint8_t* data;
int r;
int c;
};
enum class DimOrder {
CHWT, // shard memory layout: (h_per, w_per, T) — T innermost
CTHW, // shard memory layout: (T, h_per, w_per) — W innermost
};
// Scatter all shards of one component into one plane region of `out`
void scatter_component(
const std::vector<ShardView>& shards,
DimOrder dim_order,
uint8_t* out,
int T,
int plane_offset,
int plane_W,
int row_stride,
int h_per,
int w_per);
// Top-level entry: schedule Y/Cb/Cr scatters in one batch for maximum thread-pool parallelism
void planar_concat(
const std::vector<ShardView>& y_shards,
int y_h_per,
int y_w_per,
const std::vector<ShardView>& cb_shards,
int uv_h_per,
int uv_w_per,
const std::vector<ShardView>& cr_shards,
DimOrder dim_order,
int T,
int H,
int W,
int out_H,
int out_W,
uint8_t* out);
// Optional knob — set the size of the static thread pool
void set_thread_pool_size(int n_threads);
} // namespace tt_dit_planar
|