code/tt_moge/kernels/pixshuf_compute.cpp · changh95/moge-2-p150 at main

moge-2-p150 / code /tt_moge /kernels /pixshuf_compute.cpp

Optimized build (2026-10-03): model call 19.0 ms, trace 17.7 ms

13b4736 verified 25 days ago

617 Bytes

// SPDX-License-Identifier: Apache-2.0
// Pixel-shuffle front (moge-2 opt kernel, ttnn.generic_op): untilize tile rows of the [HW, 4*co] TILE matmul
// output, one dy half (WT = 2*co/32 tiles) per work unit, into c_16 (32 row-major rows). Runtime arg 0: units.
#include <cstdint>
#include "ttnn/cpp/ttnn/kernel_lib/untilize_helpers.hpp"

void kernel_main() {
    const uint32_t ntr = get_arg_val<uint32_t>(0);
    constexpr uint32_t WT = get_compile_time_arg_val(0);
    compute_kernel_hw_startup(tt::CBIndex::c_0, tt::CBIndex::c_16);
    compute_kernel_lib::untilize<WT, tt::CBIndex::c_0, tt::CBIndex::c_16>(ntr);
}

	// SPDX-License-Identifier: Apache-2.0
	// Pixel-shuffle front (moge-2 opt kernel, ttnn.generic_op): untilize tile rows of the [HW, 4*co] TILE matmul
	// output, one dy half (WT = 2*co/32 tiles) per work unit, into c_16 (32 row-major rows). Runtime arg 0: units.
	#include <cstdint>
	#include "ttnn/cpp/ttnn/kernel_lib/untilize_helpers.hpp"

	void kernel_main() {
	const uint32_t ntr = get_arg_val<uint32_t>(0);
	constexpr uint32_t WT = get_compile_time_arg_val(0);
	compute_kernel_hw_startup(tt::CBIndex::c_0, tt::CBIndex::c_16);
	compute_kernel_lib::untilize<WT, tt::CBIndex::c_0, tt::CBIndex::c_16>(ntr);
	}