From 7f9788fbf76380a26cc8ea06ac5f72191528e21b Mon Sep 17 00:00:00 2001 From: anastasios Date: Sat, 27 Jun 2026 05:47:21 +0000 Subject: [PATCH 1/4] feature(causal_conv1d): import kernel from jit_cpp --- CMakeLists.txt | 1 + csrc/host/pybind11.cpp | 6 + csrc/host/torch_causal_conv1d.h | 117 +++++++ csrc/kernel/kernel_causal_conv1d.cpp | 491 +++++++++++++++++++++++++++ 4 files changed, 615 insertions(+) create mode 100644 csrc/host/torch_causal_conv1d.h create mode 100644 csrc/kernel/kernel_causal_conv1d.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 40c2a84d..c858515d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -124,6 +124,7 @@ ascendc_library( SHARED csrc/kernel/kernel_tri_inv_col_sweep.cpp csrc/kernel/kernel_abs.cpp + csrc/kernel/kernel_causal_conv1d.cpp csrc/kernel/kernel_chunk_cumsum.cpp csrc/kernel/kernel_csr_gather.cpp csrc/kernel/kernel_gdn_chunk_h.cpp diff --git a/csrc/host/pybind11.cpp b/csrc/host/pybind11.cpp index 099c964e..d8405add 100644 --- a/csrc/host/pybind11.cpp +++ b/csrc/host/pybind11.cpp @@ -11,6 +11,7 @@ for the full License text. #include "torch_abs.h" #include "torch_batch_matrix_square.h" +#include "torch_causal_conv1d.h" #include "torch_chunk_cumsum.h" #include "torch_csr_gather.h" #include "torch_gdn_chunk_h.h" @@ -43,6 +44,11 @@ PYBIND11_MODULE(pto_kernels_ops, m) { }, pybind11::arg("device_id") = 0); m.def("pto_abs", &pto_isa_ops::run_abs); + m.def("pto_causal_conv1d", &pto_isa_ops::run_causal_conv1d, py::arg("x"), + py::arg("weights"), py::arg("bias")); + m.def("pto_causal_conv1d_batched", &pto_isa_ops::run_causal_conv1d_batched, + py::arg("x"), py::arg("weights"), py::arg("bias"), + py::arg("activation") = true); m.def("pto_chunk_h", &pto_isa_ops::run_gdn_chunk_h, py::arg("K"), py::arg("W"), py::arg("U"), py::arg("G"), py::arg("cu_seqlens") = at::zeros({1}), py::arg("batch_size"), diff --git a/csrc/host/torch_causal_conv1d.h b/csrc/host/torch_causal_conv1d.h new file mode 100644 index 00000000..372b798b --- /dev/null +++ b/csrc/host/torch_causal_conv1d.h @@ -0,0 +1,117 @@ +/** +Copyright (c) 2026 Huawei Technologies Co., Ltd. +All rights reserved. + +See LICENSE in the root of the software repository: +https://github.com/huawei-csl/pto-kernels/ +for the full License text. +*/ +#pragma once + +#include +#include + +#include "aclrtlaunch_causal_conv1d_batched_bf16_kernel.h" +#include "aclrtlaunch_causal_conv1d_batched_kernel.h" +#include "aclrtlaunch_causal_conv1d_kernel.h" +#include "utils.h" + +namespace pto_isa_ops { + +/** + * @brief Depthwise causal conv1d + per-channel bias + SiLU (single sequence). + * + * Computes y[i,c] = silu(bias[c] + sum_{k} W[k,c] * x[i-K+1+k, c]) + * with zero-padding for i < 0. Always applies SiLU activation. + * + * @param [in] x Input tensor [seqLen, channels] fp16, contiguous. + * @param [in] weights Filter weights [K, channels] fp32, contiguous. + * @param [in] bias Per-channel bias [channels] fp32, contiguous. + * @return at::Tensor Output [seqLen, channels] fp16. + */ +at::Tensor run_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, + const at::Tensor& bias) { + TORCH_CHECK(x.device().type() == DEVICE_TYPE, + "causal_conv1d: x must be on NPU, got ", x.device()); + TORCH_CHECK(x.scalar_type() == at::kHalf, + "causal_conv1d: x must be fp16, got ", x.scalar_type()); + TORCH_CHECK(x.dim() == 2, + "causal_conv1d: x must be 2D [seqLen, channels], got ", x.dim(), + "D"); + TORCH_CHECK(x.is_contiguous(), "causal_conv1d: x must be contiguous"); + TORCH_CHECK(weights.scalar_type() == at::kFloat, + "causal_conv1d: weights must be fp32, got ", + weights.scalar_type()); + TORCH_CHECK(weights.is_contiguous(), + "causal_conv1d: weights must be contiguous"); + TORCH_CHECK(bias.dim() == 1 && bias.scalar_type() == at::kFloat, + "causal_conv1d: bias must be 1D fp32"); + TORCH_CHECK(bias.is_contiguous(), "causal_conv1d: bias must be contiguous"); + + const uint32_t seqLen = static_cast(x.size(0)); + const uint32_t channels = static_cast(x.size(1)); + + at::Tensor output = at::empty_like(x); + const uint32_t block_dim = GetNumVectorCores(); + + EXEC_KERNEL_CMD(causal_conv1d_kernel, block_dim, x, output, weights, bias, + seqLen, channels); + return output; +} + +/** + * @brief Depthwise causal conv1d + per-channel bias + optional SiLU (batched). + * + * Computes y[b,i,c] = act(bias[c] + sum_{k} W[k,c] * x[b,i-K+1+k,c]) + * with zero-padding for i < 0. Supports fp16 and bf16 I/O with fp32 + * accumulation. + * + * @param [in] x Input tensor [batch, seqLen, channels] fp16 or bf16. + * @param [in] weights Filter weights [K, channels] fp32, contiguous. + * @param [in] bias Per-channel bias [channels] fp32, contiguous. + * @param [in] activation Whether to apply SiLU after bias add (default true). + * @return at::Tensor Output same shape and dtype as x. + */ +at::Tensor run_causal_conv1d_batched(const at::Tensor& x, + const at::Tensor& weights, + const at::Tensor& bias, + bool activation = true) { + TORCH_CHECK(x.device().type() == DEVICE_TYPE, + "causal_conv1d_batched: x must be on NPU, got ", x.device()); + TORCH_CHECK(x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16, + "causal_conv1d_batched: x must be fp16 or bf16, got ", + x.scalar_type()); + TORCH_CHECK(x.dim() == 3, + "causal_conv1d_batched: x must be 3D [batch, seqLen, channels], " + "got ", + x.dim(), "D"); + TORCH_CHECK(x.is_contiguous(), "causal_conv1d_batched: x must be contiguous"); + TORCH_CHECK(weights.scalar_type() == at::kFloat, + "causal_conv1d_batched: weights must be fp32, got ", + weights.scalar_type()); + TORCH_CHECK(weights.is_contiguous(), + "causal_conv1d_batched: weights must be contiguous"); + TORCH_CHECK(bias.dim() == 1 && bias.scalar_type() == at::kFloat, + "causal_conv1d_batched: bias must be 1D fp32"); + TORCH_CHECK(bias.is_contiguous(), + "causal_conv1d_batched: bias must be contiguous"); + + const uint32_t batch = static_cast(x.size(0)); + const uint32_t seqLen = static_cast(x.size(1)); + const uint32_t channels = static_cast(x.size(2)); + const uint32_t applyActivation = activation ? 1u : 0u; + + at::Tensor output = at::empty_like(x); + const uint32_t block_dim = GetNumVectorCores(); + + if (x.scalar_type() == at::kHalf) { + EXEC_KERNEL_CMD(causal_conv1d_batched_kernel, block_dim, x, output, weights, + bias, batch, seqLen, channels, applyActivation); + } else if (x.scalar_type() == at::kBFloat16) { + EXEC_KERNEL_CMD(causal_conv1d_batched_bf16_kernel, block_dim, x, output, + weights, bias, batch, seqLen, channels, applyActivation); + } + return output; +} + +} // namespace pto_isa_ops diff --git a/csrc/kernel/kernel_causal_conv1d.cpp b/csrc/kernel/kernel_causal_conv1d.cpp new file mode 100644 index 00000000..abf5ff98 --- /dev/null +++ b/csrc/kernel/kernel_causal_conv1d.cpp @@ -0,0 +1,491 @@ +/** +Copyright (c) 2026 Huawei Technologies Co., Ltd. +All rights reserved. + +See LICENSE in the root of the software repository: +https://github.com/huawei-csl/pto-kernels/ +for the full License text. +*/ + +#include "kernel_utils.h" + +using namespace pto; + +#define DIV_ROUNDUP(x, y) (((x) + (y) - 1) / (y)) +#define ALIGN_UP(x, y) (DIV_ROUNDUP((x), (y)) * (y)) + +// =========================================================================== +// DEPTHWISE fused causal conv1d + bias + (optional) SiLU (per-channel, any K) +// +// y[b,i,c] = act( bias[c] + sum_{k=max(0,K-1-i)..K-1} W[k,c] * x[b, i-K+1+k, +// c] ), x[<0]=0 +// +// Per-channel K-tap filter (Mamba/GDN short conv). x,y are [batch, seqLen, +// channels] row-major; `channels` is the lane axis, seqLen the conv axis. +// Weights W[K,channels] + bias[channels] are fp32 GM tensors. fp16 OR bf16 I/O, +// fp32 accumulate. +// +// Filter width K and per-tile channel width MAX_W are compile-time constants +// chosen at the call site as template parameters (no preprocessor config), e.g. +// constexpr uint32_t K = CAUSAL_CONV_K, MAX_W = CAUSAL_CONV_MAX_W; +// csilu::runConvSiluBatched(...); +// +// 2-D-plus-batch work grid: workUnits = batch x sequenceChunkCount x +// channelTileCount. Each work unit produces outputs [batchIndex] x +// [outputRowStart,outputRowEnd) for channels +// [channelTileBase,channelTileBase+tileChannelCount), replaying K-1 causal halo +// rows to prime its accumulators. The grid fills all cores: batch supplies +// parallelism first, then channel tiles, then sequence chunks; the channel-tile +// width is widened to whole vector lanes for coalesced stores. (See +// processWorkUnit + runConvSiluBatched.) +// +// Generic in K (filter width) and MAX_W (tile width). accumRingSize = smallest +// power of two >= K is the accumulator ring size, so the K outputs in flight +// map to distinct ring slots via `idx & accumRingMask` (accumRingMask = +// accumRingSize - 1). +// UB layout (per lane, fp32 unless noted): K weights + bias(1) + accumRingSize +// accumulators + (K-1) partial products + input-as-fp32(1) = +// 2*K+accumRingSize+1 fp32 tiles; then the I/O region inputTile[0..1] + output0 +// + output1 = 4 I/O tiles (input load double-buffered). A static_assert keeps +// the total within UB_BYTES_PER_CORE for the chosen K / MAX_W / dtypes. NOTE: +// uses the PTO tile-op API (); the `csilu` namespace avoids a +// clash with pto::detail. +// =========================================================================== + +namespace csilu { + +// Unified Buffer available per AIV core (Ascend 910B2 = 192 KiB). The UB +// static_assert in processWorkUnit checks the chosen layout fits; raise it for +// a next-gen NPU with a larger UB. +constexpr uint32_t UB_BYTES_PER_CORE = 192u * 1024u; + +// Smallest power of two >= value. +AICORE constexpr uint32_t roundUpToPowerOfTwo(uint32_t value) { + if (value != 0u) --value; + + value |= (value >> 1u); + value |= (value >> 2u); + value |= (value >> 4u); + value |= (value >> 8u); + value |= (value >> 16u); + + ++value; + + return value; +} + +template +AICORE inline void applySiluToTile(TileT& dst, TileT& src, TileT& scratch) { + using ElemType = typename TileT::DType; + TMULS(scratch, src, (ElemType)-1); + pipe_barrier(PIPE_V); + TEXP(scratch, scratch); + pipe_barrier(PIPE_V); + TADDS(scratch, scratch, (ElemType)1); + pipe_barrier(PIPE_V); + TDIV(dst, src, scratch); +} + +// Process ONE work unit: outputs [outputRowStart,outputRowEnd) for channels +// [channelTileBase,channelTileBase+tileChannelCount) of the sequence whose +// first row is at element offset sequenceRowOffset. x[<0]=0 (no cache). +template +AICORE inline void processWorkUnit( + __gm__ IoElemType* input, __gm__ IoElemType* output, + __gm__ AccumElemType* weights, __gm__ AccumElemType* bias, + uint32_t channels, uint64_t sequenceRowOffset, uint32_t channelTileBase, + int32_t tileChannelCount, uint32_t outputRowStart, uint32_t outputRowEnd, + uint32_t applyActivation) { + using GlobalShape = pto::Shape<1, 1, 1, 1, DYNAMIC>; + using GlobalStride = pto::Stride<1, 1, 1, 1, 1>; + using GlobalIoTensor = + pto::GlobalTensor; + using GlobalAccumTensor = + pto::GlobalTensor; + using IoTile = + Tile; + using AccumTile = Tile; + + constexpr uint32_t accumTileBytes = MAX_W * sizeof(AccumElemType); + constexpr uint32_t ioTileBytes = MAX_W * sizeof(IoElemType); + // accumulator ring (power of two >= K) so the K in-flight outputs never + // alias. + constexpr uint32_t accumRingSize = roundUpToPowerOfTwo(K); + constexpr uint32_t accumRingMask = + accumRingSize - 1u; // ring-slot index mask + static_assert(K <= accumRingSize, "accumulator ring must hold all K taps"); + + // UB byte offsets. fp32 region: K weights (weight k at k*accumTileBytes) | + // bias | accumRingSize accumulators | K-1 partial products | input-as-fp32. + // Then the I/O region: 4 ioTileBytes-sized tiles (input load + // double-buffered). + constexpr uint32_t ubBiasOffset = K * accumTileBytes; + constexpr uint32_t ubAccumRingBase = (K + 1u) * accumTileBytes; + // partial product for tap k at ubProductBase + (k-1)*accumTileBytes; also + // reused as the SiLU scratch tile once the products have been summed. + constexpr uint32_t ubProductBase = (K + 1u + accumRingSize) * accumTileBytes; + constexpr uint32_t ubInputFp32Offset = + (2u * K + accumRingSize) * accumTileBytes; + constexpr uint32_t ubIoRegionBase = + (2u * K + accumRingSize + 1u) * accumTileBytes; + static_assert( + ubIoRegionBase + 4u * ioTileBytes <= UB_BYTES_PER_CORE, + "conv1d UB exceeds UB_BYTES_PER_CORE: lower K/MAX_W or raise it"); + + constexpr uint32_t ubOutputOffset[2] = {ubIoRegionBase + ioTileBytes, + ubIoRegionBase + 2u * ioTileBytes}; + // input double-buffer: slot 0 before the outputs, slot 1 after. + constexpr uint32_t ubInputOffset[2] = {ubIoRegionBase, + ubIoRegionBase + 3u * ioTileBytes}; + + const uint32_t firstInputRow = + (outputRowStart > (K - 1)) ? (outputRowStart - (K - 1)) : 0u; + + // ---- per-channel weights + bias (resident for this work unit) ---- + for (uint32_t tapIndex = 0; tapIndex < K; ++tapIndex) { + GlobalAccumTensor weightGm( + weights + (uint64_t)tapIndex * channels + channelTileBase, + {tileChannelCount}); + AccumTile weightTile(tileChannelCount); + TASSIGN(weightTile, tapIndex * accumTileBytes); + TLOAD(weightTile, weightGm); + } + { + GlobalAccumTensor biasGm(bias + channelTileBase, {tileChannelCount}); + AccumTile biasTile(tileChannelCount); + TASSIGN(biasTile, ubBiasOffset); + TLOAD(biasTile, biasGm); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + + // double-buffered input: two load slots with independent handshakes. + // EVENT_ID3 is reused here (the weight load above already consumed it). + // events are passed by mutable value to set_flag/wait_flag, so not constexpr. + const event_t inputBufferEvent[2] = {EVENT_ID0, EVENT_ID3}; + set_flag(PIPE_V, PIPE_MTE2, + inputBufferEvent[0]); // input slot 0 initially free + set_flag(PIPE_V, PIPE_MTE2, + inputBufferEvent[1]); // input slot 1 initially free + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); + + // PROLOGUE: issue the first input load (firstInputRow) so iteration 0 can + // prefetch the next row. + if (firstInputRow < outputRowEnd) { + GlobalIoTensor inputRowGm(input + sequenceRowOffset + + (uint64_t)firstInputRow * channels + + channelTileBase, + {tileChannelCount}); + IoTile prologueInputTile(tileChannelCount); + TASSIGN(prologueInputTile, ubInputOffset[0]); + wait_flag(PIPE_V, PIPE_MTE2, inputBufferEvent[0]); + TLOAD(prologueInputTile, inputRowGm); + set_flag(PIPE_MTE2, PIPE_V, inputBufferEvent[0]); + } + + for (uint32_t inputRow = firstInputRow; inputRow < outputRowEnd; ++inputRow) { + const uint32_t bufferIndex = (inputRow - firstInputRow) & 1u; + IoTile inputTileIo(tileChannelCount); + AccumTile inputTileFp32(tileChannelCount); + TASSIGN(inputTileIo, ubInputOffset[bufferIndex]); + TASSIGN(inputTileFp32, ubInputFp32Offset); + + // (1) consume current row from buffer bufferIndex (loaded by the prologue / + // previous prefetch). + wait_flag(PIPE_MTE2, PIPE_V, inputBufferEvent[bufferIndex]); + TCVT(inputTileFp32, inputTileIo, pto::RoundMode::CAST_NONE); + set_flag(PIPE_V, PIPE_MTE2, + inputBufferEvent[bufferIndex]); // slot free again + + // (2) prefetch next row into the OTHER buffer; overlaps the compute below + // without waiting on the buffer consumed this iteration. + if (inputRow + 1 < outputRowEnd) { + const uint32_t nextBufferIndex = bufferIndex ^ 1u; + IoTile nextInputTile(tileChannelCount); + TASSIGN(nextInputTile, ubInputOffset[nextBufferIndex]); + GlobalIoTensor nextInputRowGm(input + sequenceRowOffset + + (uint64_t)(inputRow + 1) * channels + + channelTileBase, + {tileChannelCount}); + wait_flag(PIPE_V, PIPE_MTE2, inputBufferEvent[nextBufferIndex]); + TLOAD(nextInputTile, nextInputRowGm); + set_flag(PIPE_MTE2, PIPE_V, inputBufferEvent[nextBufferIndex]); + } + + pipe_barrier(PIPE_V); + + // scatter: this input row contributes to K output rows; form each + // weight*input product (only for outputs in [outputRowStart,outputRowEnd)). + for (uint32_t tapIndex = 0; tapIndex < K; ++tapIndex) { + const uint32_t outputRow = inputRow + (K - 1) - tapIndex; + if (outputRow < outputRowStart || outputRow >= outputRowEnd) continue; + AccumTile weightTile(tileChannelCount); + TASSIGN(weightTile, tapIndex * accumTileBytes); + if (inputRow == 0 || tapIndex == 0) { + // first contribution to this output's accumulator slot: initialise it. + AccumTile accumTile(tileChannelCount); + TASSIGN(accumTile, + ubAccumRingBase + (outputRow & accumRingMask) * accumTileBytes); + TMUL(accumTile, inputTileFp32, weightTile); + } else { + AccumTile productTile(tileChannelCount); + TASSIGN(productTile, ubProductBase + (tapIndex - 1u) * accumTileBytes); + TMUL(productTile, inputTileFp32, weightTile); + } + } + pipe_barrier(PIPE_V); + if (inputRow != 0) { + for (uint32_t tapIndex = 1; tapIndex < K; ++tapIndex) { + const uint32_t outputRow = inputRow + (K - 1) - tapIndex; + if (outputRow < outputRowStart || outputRow >= outputRowEnd) continue; + AccumTile accumTile(tileChannelCount); + AccumTile productTile(tileChannelCount); + TASSIGN(accumTile, + ubAccumRingBase + (outputRow & accumRingMask) * accumTileBytes); + TASSIGN(productTile, ubProductBase + (tapIndex - 1u) * accumTileBytes); + TADD(accumTile, accumTile, productTile); + } + } + pipe_barrier(PIPE_V); + + if (inputRow < outputRowStart) + continue; // halo row: primed accumulators only + + const uint32_t accumRingSlot = inputRow & accumRingMask; + const uint32_t outputBufferIndex = inputRow & 1u; + const event_t outputBufferEvent = (event_t)(1u + outputBufferIndex); + AccumTile accumTile(tileChannelCount); + AccumTile biasTile(tileChannelCount); + AccumTile siluScratchTile(tileChannelCount); + IoTile outputTile(tileChannelCount); + TASSIGN(accumTile, ubAccumRingBase + accumRingSlot * accumTileBytes); + TASSIGN(biasTile, ubBiasOffset); + TASSIGN(siluScratchTile, ubProductBase); + TASSIGN(outputTile, ubOutputOffset[outputBufferIndex]); + + TADD(accumTile, accumTile, biasTile); + pipe_barrier(PIPE_V); + if (applyActivation) { + applySiluToTile(accumTile, accumTile, siluScratchTile); + pipe_barrier(PIPE_V); + } + wait_flag(PIPE_MTE3, PIPE_V, outputBufferEvent); + TCVT(outputTile, accumTile, pto::RoundMode::CAST_NONE); + + GlobalIoTensor outputRowGm(output + sequenceRowOffset + + (uint64_t)inputRow * channels + + channelTileBase, + {tileChannelCount}); + set_flag(PIPE_V, PIPE_MTE3, outputBufferEvent); + wait_flag(PIPE_V, PIPE_MTE3, outputBufferEvent); + TSTORE(outputRowGm, outputTile); + set_flag(PIPE_MTE3, PIPE_V, outputBufferEvent); + } + + wait_flag(PIPE_V, PIPE_MTE2, inputBufferEvent[0]); + wait_flag(PIPE_V, PIPE_MTE2, inputBufferEvent[1]); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); +} + +template +AICORE void runConvSiluBatched(__gm__ IoElemType* input, + __gm__ IoElemType* output, + __gm__ AccumElemType* weights, + __gm__ AccumElemType* bias, uint32_t batch, + uint32_t seqLen, uint32_t channels, + uint32_t applyActivation) { + static_assert(K >= 1u, "K (filter width) must be >= 1"); + + set_mask_norm(); + set_vector_mask(-1, -1); + + const uint32_t coreCount = get_block_num(); + const uint32_t coreIndex = get_block_idx(); + if (seqLen == 0 || batch == 0 || channels == 0) return; + + // ---- grid tuning knobs ---- + // Minimum rows per sequence-chunk. Smaller chunks expose more parallelism but + // each replays K-1 causal halo rows, so very small chunks waste compute; 32 + // is the balance point across the GDN sequence lengths. + constexpr uint32_t minSeqChunkLen = 32u; + // Channels per aligned vector lane = 256 B / element size (128 for + // fp16/bf16). Channel tiles are sized in whole lanes for coalesced + // loads/stores. + constexpr uint32_t channelsPerLane = 256u / sizeof(IoElemType); + + // batch supplies parallelism first; each sequence then needs + // workUnitsPerSequence (channel-tile x seq-chunk) units to keep all cores + // busy. + uint32_t workUnitsPerSequence = DIV_ROUNDUP(coreCount, batch); + if (workUnitsPerSequence < 1) workUnitsPerSequence = 1; + + uint32_t maxSequenceChunks = DIV_ROUNDUP(seqLen, minSeqChunkLen); + if (maxSequenceChunks < 1) maxSequenceChunks = 1; + + // channel tiles: enough that each tile fits MAX_W, and enough (with the chunk + // count) to fill the per-sequence work-unit budget. + const uint32_t channelTilesForUbLimit = DIV_ROUNDUP(channels, MAX_W); + const uint32_t channelTilesToFillCores = + DIV_ROUNDUP(workUnitsPerSequence, maxSequenceChunks); + uint32_t channelTileCount = channelTilesForUbLimit > channelTilesToFillCores + ? channelTilesForUbLimit + : channelTilesToFillCores; + + const uint32_t maxChannelTilesByLane = DIV_ROUNDUP(channels, channelsPerLane); + if (channelTileCount > maxChannelTilesByLane) + channelTileCount = maxChannelTilesByLane; + if (channelTileCount < 1) channelTileCount = 1; + + uint32_t channelTileWidth = + ALIGN_UP(DIV_ROUNDUP(channels, channelTileCount), channelsPerLane); + if (channelTileWidth < channelsPerLane) channelTileWidth = channelsPerLane; + if (channelTileWidth > MAX_W) channelTileWidth = MAX_W; + if (channelTileWidth > channels) channelTileWidth = channels; + + channelTileCount = DIV_ROUNDUP(channels, channelTileWidth); + uint32_t sequenceChunkCount = + DIV_ROUNDUP(workUnitsPerSequence, channelTileCount); + if (sequenceChunkCount < 1) sequenceChunkCount = 1; + if (sequenceChunkCount > maxSequenceChunks) + sequenceChunkCount = maxSequenceChunks; + + const uint32_t sequenceChunkLength = DIV_ROUNDUP(seqLen, sequenceChunkCount); + + // Iterate the convolution (sequence) direction in the middle so input rows + // and weights stay resident across chunks of the same channel tile. + const uint32_t totalWorkUnits = batch * channelTileCount * sequenceChunkCount; + for (uint32_t workUnitIndex = coreIndex; workUnitIndex < totalWorkUnits; + workUnitIndex += coreCount) { + const uint32_t sequenceChunkIndex = workUnitIndex % sequenceChunkCount; + const uint32_t tileBatchIndex = workUnitIndex / sequenceChunkCount; + const uint32_t channelTileIndex = tileBatchIndex % channelTileCount; + const uint32_t batchIndex = tileBatchIndex / channelTileCount; + const uint32_t channelTileBase = channelTileIndex * channelTileWidth; + const uint32_t remainingChannels = channels - channelTileBase; + const int32_t tileChannelCount = remainingChannels > channelTileWidth + ? (int32_t)channelTileWidth + : (int32_t)remainingChannels; + const uint32_t outputRowStart = sequenceChunkIndex * sequenceChunkLength; + if (outputRowStart >= seqLen) continue; + uint32_t outputRowEnd = outputRowStart + sequenceChunkLength; + if (outputRowEnd > seqLen) outputRowEnd = seqLen; + const uint64_t sequenceRowOffset = (uint64_t)batchIndex * seqLen * channels; + processWorkUnit( + input, output, weights, bias, channels, sequenceRowOffset, + channelTileBase, tileChannelCount, outputRowStart, outputRowEnd, + applyActivation); + } +} + +} // namespace csilu + +// Filter width / per-tile channel width the entry points below are compiled at. +// Default to the K=4, MAX_W=3072 production configuration; the test suite +// recompiles this file at other widths via -DCAUSAL_CONV_K / +// -DCAUSAL_CONV_MAX_W (see test_causal_conv1d.py), so a plain build is +// byte-for-byte unaffected. +#ifndef CAUSAL_CONV_K +#define CAUSAL_CONV_K 4 +#endif +#ifndef CAUSAL_CONV_MAX_W +#define CAUSAL_CONV_MAX_W 3072 +#endif + +// ---- single-sequence entry (back-compat: input,output [seqLen,channels] fp16, +// weights[K,channels]/bias[channels] fp32) ---- +extern "C" __global__ AICORE void causal_conv1d_kernel( + __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, + __gm__ uint8_t* bias, uint32_t seqLen, uint32_t channels) { +#if defined(__DAV_VEC__) + constexpr uint32_t K = CAUSAL_CONV_K, MAX_W = CAUSAL_CONV_MAX_W; + csilu::runConvSiluBatched( + (__gm__ half*)input, (__gm__ half*)output, (__gm__ float*)weights, + (__gm__ float*)bias, 1u, seqLen, channels, 1u); +#else + (void)input; + (void)output; + (void)weights; + (void)bias; + (void)seqLen; + (void)channels; +#endif +} + +// ---- batched fp16 entry: input,output [batch,seqLen,channels] fp16, +// weights/bias fp32 ---- +extern "C" __global__ AICORE void causal_conv1d_batched_kernel( + __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, + __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, + uint32_t applyActivation) { +#if defined(__DAV_VEC__) + constexpr uint32_t K = CAUSAL_CONV_K, MAX_W = CAUSAL_CONV_MAX_W; + csilu::runConvSiluBatched( + (__gm__ half*)input, (__gm__ half*)output, (__gm__ float*)weights, + (__gm__ float*)bias, batch, seqLen, channels, applyActivation); +#else + (void)input; + (void)output; + (void)weights; + (void)bias; + (void)batch; + (void)seqLen; + (void)channels; + (void)applyActivation; +#endif +} + +// ---- batched bf16 entry: input,output [batch,seqLen,channels] bf16, +// weights/bias fp32 ---- +extern "C" __global__ AICORE void causal_conv1d_batched_bf16_kernel( + __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, + __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, + uint32_t applyActivation) { +#if defined(__DAV_VEC__) + constexpr uint32_t K = CAUSAL_CONV_K, MAX_W = CAUSAL_CONV_MAX_W; + csilu::runConvSiluBatched( + (__gm__ bfloat16_t*)input, (__gm__ bfloat16_t*)output, + (__gm__ float*)weights, (__gm__ float*)bias, batch, seqLen, channels, + applyActivation); +#else + (void)input; + (void)output; + (void)weights; + (void)bias; + (void)batch; + (void)seqLen; + (void)channels; + (void)applyActivation; +#endif +} + +extern "C" void call_kernel(uint32_t blockDim, void* stream, uint8_t* input, + uint8_t* output, uint8_t* weights, uint8_t* bias, + uint32_t seqLen, uint32_t channels) { + causal_conv1d_kernel<<>>( + input, output, weights, bias, seqLen, channels); +} + +extern "C" void call_kernel_batched(uint32_t blockDim, void* stream, + uint8_t* input, uint8_t* output, + uint8_t* weights, uint8_t* bias, + uint32_t batch, uint32_t seqLen, + uint32_t channels, + uint32_t applyActivation) { + causal_conv1d_batched_kernel<<>>( + input, output, weights, bias, batch, seqLen, channels, applyActivation); +} + +extern "C" void call_kernel_batched_bf16(uint32_t blockDim, void* stream, + uint8_t* input, uint8_t* output, + uint8_t* weights, uint8_t* bias, + uint32_t batch, uint32_t seqLen, + uint32_t channels, + uint32_t applyActivation) { + causal_conv1d_batched_bf16_kernel<<>>( + input, output, weights, bias, batch, seqLen, channels, applyActivation); +} From 1c6354e5a669e5e482ebe12369d6404c3e833aab Mon Sep 17 00:00:00 2001 From: anastasios Date: Sat, 27 Jun 2026 06:07:10 +0000 Subject: [PATCH 2/4] fix --- CMakeLists.txt | 2 +- csrc/host/pybind11.cpp | 12 +- ...sal_conv1d.h => torch_gdn_causal_conv1d.h} | 69 +++++----- ...onv1d.cpp => kernel_gdn_causal_conv1d.cpp} | 122 ++++++++++++------ tests/test_gdn_causal_conv1d.py | 87 +++++++++++++ 5 files changed, 215 insertions(+), 77 deletions(-) rename csrc/host/{torch_causal_conv1d.h => torch_gdn_causal_conv1d.h} (58%) rename csrc/kernel/{kernel_causal_conv1d.cpp => kernel_gdn_causal_conv1d.cpp} (82%) create mode 100644 tests/test_gdn_causal_conv1d.py diff --git a/CMakeLists.txt b/CMakeLists.txt index c858515d..965b5bc4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -124,7 +124,7 @@ ascendc_library( SHARED csrc/kernel/kernel_tri_inv_col_sweep.cpp csrc/kernel/kernel_abs.cpp - csrc/kernel/kernel_causal_conv1d.cpp + csrc/kernel/kernel_gdn_causal_conv1d.cpp csrc/kernel/kernel_chunk_cumsum.cpp csrc/kernel/kernel_csr_gather.cpp csrc/kernel/kernel_gdn_chunk_h.cpp diff --git a/csrc/host/pybind11.cpp b/csrc/host/pybind11.cpp index d8405add..9d958dc9 100644 --- a/csrc/host/pybind11.cpp +++ b/csrc/host/pybind11.cpp @@ -11,9 +11,9 @@ for the full License text. #include "torch_abs.h" #include "torch_batch_matrix_square.h" -#include "torch_causal_conv1d.h" #include "torch_chunk_cumsum.h" #include "torch_csr_gather.h" +#include "torch_gdn_causal_conv1d.h" #include "torch_gdn_chunk_h.h" #include "torch_gdn_chunk_o.h" #include "torch_gdn_scaled_dot_kkt.h" @@ -44,11 +44,11 @@ PYBIND11_MODULE(pto_kernels_ops, m) { }, pybind11::arg("device_id") = 0); m.def("pto_abs", &pto_isa_ops::run_abs); - m.def("pto_causal_conv1d", &pto_isa_ops::run_causal_conv1d, py::arg("x"), - py::arg("weights"), py::arg("bias")); - m.def("pto_causal_conv1d_batched", &pto_isa_ops::run_causal_conv1d_batched, - py::arg("x"), py::arg("weights"), py::arg("bias"), - py::arg("activation") = true); + m.def("pto_gdn_causal_conv1d", &pto_isa_ops::run_gdn_causal_conv1d, + py::arg("x"), py::arg("weights"), py::arg("bias")); + m.def("pto_gdn_causal_conv1d_batched", + &pto_isa_ops::run_gdn_causal_conv1d_batched, py::arg("x"), + py::arg("weights"), py::arg("bias"), py::arg("activation") = true); m.def("pto_chunk_h", &pto_isa_ops::run_gdn_chunk_h, py::arg("K"), py::arg("W"), py::arg("U"), py::arg("G"), py::arg("cu_seqlens") = at::zeros({1}), py::arg("batch_size"), diff --git a/csrc/host/torch_causal_conv1d.h b/csrc/host/torch_gdn_causal_conv1d.h similarity index 58% rename from csrc/host/torch_causal_conv1d.h rename to csrc/host/torch_gdn_causal_conv1d.h index 372b798b..00153d0c 100644 --- a/csrc/host/torch_causal_conv1d.h +++ b/csrc/host/torch_gdn_causal_conv1d.h @@ -11,9 +11,9 @@ for the full License text. #include #include -#include "aclrtlaunch_causal_conv1d_batched_bf16_kernel.h" -#include "aclrtlaunch_causal_conv1d_batched_kernel.h" -#include "aclrtlaunch_causal_conv1d_kernel.h" +#include "aclrtlaunch_gdn_causal_conv1d_batched_bf16_kernel.h" +#include "aclrtlaunch_gdn_causal_conv1d_batched_kernel.h" +#include "aclrtlaunch_gdn_causal_conv1d_kernel.h" #include "utils.h" namespace pto_isa_ops { @@ -29,24 +29,25 @@ namespace pto_isa_ops { * @param [in] bias Per-channel bias [channels] fp32, contiguous. * @return at::Tensor Output [seqLen, channels] fp16. */ -at::Tensor run_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, - const at::Tensor& bias) { +at::Tensor run_gdn_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, + const at::Tensor& bias) { TORCH_CHECK(x.device().type() == DEVICE_TYPE, - "causal_conv1d: x must be on NPU, got ", x.device()); + "gdn_causal_conv1d: x must be on NPU, got ", x.device()); TORCH_CHECK(x.scalar_type() == at::kHalf, - "causal_conv1d: x must be fp16, got ", x.scalar_type()); + "gdn_causal_conv1d: x must be fp16, got ", x.scalar_type()); TORCH_CHECK(x.dim() == 2, - "causal_conv1d: x must be 2D [seqLen, channels], got ", x.dim(), - "D"); - TORCH_CHECK(x.is_contiguous(), "causal_conv1d: x must be contiguous"); + "gdn_causal_conv1d: x must be 2D [seqLen, channels], got ", + x.dim(), "D"); + TORCH_CHECK(x.is_contiguous(), "gdn_causal_conv1d: x must be contiguous"); TORCH_CHECK(weights.scalar_type() == at::kFloat, - "causal_conv1d: weights must be fp32, got ", + "gdn_causal_conv1d: weights must be fp32, got ", weights.scalar_type()); TORCH_CHECK(weights.is_contiguous(), - "causal_conv1d: weights must be contiguous"); + "gdn_causal_conv1d: weights must be contiguous"); TORCH_CHECK(bias.dim() == 1 && bias.scalar_type() == at::kFloat, - "causal_conv1d: bias must be 1D fp32"); - TORCH_CHECK(bias.is_contiguous(), "causal_conv1d: bias must be contiguous"); + "gdn_causal_conv1d: bias must be 1D fp32"); + TORCH_CHECK(bias.is_contiguous(), + "gdn_causal_conv1d: bias must be contiguous"); const uint32_t seqLen = static_cast(x.size(0)); const uint32_t channels = static_cast(x.size(1)); @@ -54,7 +55,7 @@ at::Tensor run_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, at::Tensor output = at::empty_like(x); const uint32_t block_dim = GetNumVectorCores(); - EXEC_KERNEL_CMD(causal_conv1d_kernel, block_dim, x, output, weights, bias, + EXEC_KERNEL_CMD(gdn_causal_conv1d_kernel, block_dim, x, output, weights, bias, seqLen, channels); return output; } @@ -72,29 +73,31 @@ at::Tensor run_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, * @param [in] activation Whether to apply SiLU after bias add (default true). * @return at::Tensor Output same shape and dtype as x. */ -at::Tensor run_causal_conv1d_batched(const at::Tensor& x, - const at::Tensor& weights, - const at::Tensor& bias, - bool activation = true) { +at::Tensor run_gdn_causal_conv1d_batched(const at::Tensor& x, + const at::Tensor& weights, + const at::Tensor& bias, + bool activation = true) { TORCH_CHECK(x.device().type() == DEVICE_TYPE, - "causal_conv1d_batched: x must be on NPU, got ", x.device()); + "gdn_causal_conv1d_batched: x must be on NPU, got ", x.device()); TORCH_CHECK(x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16, - "causal_conv1d_batched: x must be fp16 or bf16, got ", + "gdn_causal_conv1d_batched: x must be fp16 or bf16, got ", x.scalar_type()); - TORCH_CHECK(x.dim() == 3, - "causal_conv1d_batched: x must be 3D [batch, seqLen, channels], " - "got ", - x.dim(), "D"); - TORCH_CHECK(x.is_contiguous(), "causal_conv1d_batched: x must be contiguous"); + TORCH_CHECK( + x.dim() == 3, + "gdn_causal_conv1d_batched: x must be 3D [batch, seqLen, channels], " + "got ", + x.dim(), "D"); + TORCH_CHECK(x.is_contiguous(), + "gdn_causal_conv1d_batched: x must be contiguous"); TORCH_CHECK(weights.scalar_type() == at::kFloat, - "causal_conv1d_batched: weights must be fp32, got ", + "gdn_causal_conv1d_batched: weights must be fp32, got ", weights.scalar_type()); TORCH_CHECK(weights.is_contiguous(), - "causal_conv1d_batched: weights must be contiguous"); + "gdn_causal_conv1d_batched: weights must be contiguous"); TORCH_CHECK(bias.dim() == 1 && bias.scalar_type() == at::kFloat, - "causal_conv1d_batched: bias must be 1D fp32"); + "gdn_causal_conv1d_batched: bias must be 1D fp32"); TORCH_CHECK(bias.is_contiguous(), - "causal_conv1d_batched: bias must be contiguous"); + "gdn_causal_conv1d_batched: bias must be contiguous"); const uint32_t batch = static_cast(x.size(0)); const uint32_t seqLen = static_cast(x.size(1)); @@ -105,10 +108,10 @@ at::Tensor run_causal_conv1d_batched(const at::Tensor& x, const uint32_t block_dim = GetNumVectorCores(); if (x.scalar_type() == at::kHalf) { - EXEC_KERNEL_CMD(causal_conv1d_batched_kernel, block_dim, x, output, weights, - bias, batch, seqLen, channels, applyActivation); + EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_kernel, block_dim, x, output, + weights, bias, batch, seqLen, channels, applyActivation); } else if (x.scalar_type() == at::kBFloat16) { - EXEC_KERNEL_CMD(causal_conv1d_batched_bf16_kernel, block_dim, x, output, + EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_bf16_kernel, block_dim, x, output, weights, bias, batch, seqLen, channels, applyActivation); } return output; diff --git a/csrc/kernel/kernel_causal_conv1d.cpp b/csrc/kernel/kernel_gdn_causal_conv1d.cpp similarity index 82% rename from csrc/kernel/kernel_causal_conv1d.cpp rename to csrc/kernel/kernel_gdn_causal_conv1d.cpp index abf5ff98..ce5a47c5 100644 --- a/csrc/kernel/kernel_causal_conv1d.cpp +++ b/csrc/kernel/kernel_gdn_causal_conv1d.cpp @@ -387,7 +387,7 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, // Filter width / per-tile channel width the entry points below are compiled at. // Default to the K=4, MAX_W=3072 production configuration; the test suite // recompiles this file at other widths via -DCAUSAL_CONV_K / -// -DCAUSAL_CONV_MAX_W (see test_causal_conv1d.py), so a plain build is +// -DCAUSAL_CONV_MAX_W (see test_gdn_causal_conv1d.py), so a plain build is // byte-for-byte unaffected. #ifndef CAUSAL_CONV_K #define CAUSAL_CONV_K 4 @@ -396,9 +396,34 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, #define CAUSAL_CONV_MAX_W 3072 #endif -// ---- single-sequence entry (back-compat: input,output [seqLen,channels] fp16, -// weights[K,channels]/bias[channels] fp32) ---- -extern "C" __global__ AICORE void causal_conv1d_kernel( +/** + * @brief Single-sequence causal depthwise conv1d with optional SiLU activation + * (fp16 I/O). + * + * Back-compat entry point for single-batch inference. Activations and weights + * are processed with a filter width of @c CAUSAL_CONV_K and a maximum sequence + * length of + * @c CAUSAL_CONV_MAX_W (both set at compile time; defaults: K=4, MAX_W=3072). + * + * Equivalent to calling @c gdn_causal_conv1d_batched_kernel with @c batch=1 and + * @c applyActivation=1. + * + * @param[in] input Global-memory pointer to the input tensor [seqLen, + * channels] in fp16. + * @param[out] output Global-memory pointer to the output tensor [seqLen, + * channels] in fp16. + * @param[in] weights Global-memory pointer to the convolution weights [K, + * channels] in fp32. + * @param[in] bias Global-memory pointer to the bias vector [channels] in + * fp32. + * @param[in] seqLen Number of time steps in the sequence. + * @param[in] channels Number of channels (must be a multiple of the tile + * width). + * + * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On + * other targets all parameters are ignored and the kernel is a no-op. + */ +extern "C" __global__ AICORE void gdn_causal_conv1d_kernel( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t seqLen, uint32_t channels) { #if defined(__DAV_VEC__) @@ -416,9 +441,33 @@ extern "C" __global__ AICORE void causal_conv1d_kernel( #endif } -// ---- batched fp16 entry: input,output [batch,seqLen,channels] fp16, -// weights/bias fp32 ---- -extern "C" __global__ AICORE void causal_conv1d_batched_kernel( +/** + * @brief Batched causal depthwise conv1d with optional SiLU activation (fp16 + * I/O). + * + * Processes a batch of sequences in a single kernel launch. Filter width and + * maximum sequence length are fixed at compile time via @c CAUSAL_CONV_K and @c + * CAUSAL_CONV_MAX_W (defaults: K=4, MAX_W=3072). + * + * @param[in] input Global-memory pointer to the input tensor + * [batch, seqLen, channels] in fp16. + * @param[out] output Global-memory pointer to the output tensor + * [batch, seqLen, channels] in fp16. + * @param[in] weights Global-memory pointer to the convolution weights + * [K, channels] in fp32. + * @param[in] bias Global-memory pointer to the bias vector + * [channels] in fp32. + * @param[in] batch Number of sequences in the batch. + * @param[in] seqLen Number of time steps per sequence. + * @param[in] channels Number of channels (must be a multiple of the + * tile width). + * @param[in] applyActivation Non-zero to apply SiLU after the convolution; + * zero to skip. + * + * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On + * other targets all parameters are ignored and the kernel is a no-op. + */ +extern "C" __global__ AICORE void gdn_causal_conv1d_batched_kernel( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, uint32_t applyActivation) { @@ -439,9 +488,35 @@ extern "C" __global__ AICORE void causal_conv1d_batched_kernel( #endif } -// ---- batched bf16 entry: input,output [batch,seqLen,channels] bf16, -// weights/bias fp32 ---- -extern "C" __global__ AICORE void causal_conv1d_batched_bf16_kernel( +/** + * @brief Batched causal depthwise conv1d with optional SiLU activation + * (bfloat16 I/O). + * + * Identical in semantics to @c gdn_causal_conv1d_batched_kernel but uses + * bfloat16 for input and output tensors instead of fp16. Weights and bias + * remain fp32. Filter width and maximum sequence length are fixed at compile + * time via @c CAUSAL_CONV_K and + * @c CAUSAL_CONV_MAX_W (defaults: K=4, MAX_W=3072). + * + * @param[in] input Global-memory pointer to the input tensor + * [batch, seqLen, channels] in bfloat16. + * @param[out] output Global-memory pointer to the output tensor + * [batch, seqLen, channels] in bfloat16. + * @param[in] weights Global-memory pointer to the convolution weights + * [K, channels] in fp32. + * @param[in] bias Global-memory pointer to the bias vector + * [channels] in fp32. + * @param[in] batch Number of sequences in the batch. + * @param[in] seqLen Number of time steps per sequence. + * @param[in] channels Number of channels (must be a multiple of the + * tile width). + * @param[in] applyActivation Non-zero to apply SiLU after the convolution; + * zero to skip. + * + * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On + * other targets all parameters are ignored and the kernel is a no-op. + */ +extern "C" __global__ AICORE void gdn_causal_conv1d_batched_bf16_kernel( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, uint32_t applyActivation) { @@ -462,30 +537,3 @@ extern "C" __global__ AICORE void causal_conv1d_batched_bf16_kernel( (void)applyActivation; #endif } - -extern "C" void call_kernel(uint32_t blockDim, void* stream, uint8_t* input, - uint8_t* output, uint8_t* weights, uint8_t* bias, - uint32_t seqLen, uint32_t channels) { - causal_conv1d_kernel<<>>( - input, output, weights, bias, seqLen, channels); -} - -extern "C" void call_kernel_batched(uint32_t blockDim, void* stream, - uint8_t* input, uint8_t* output, - uint8_t* weights, uint8_t* bias, - uint32_t batch, uint32_t seqLen, - uint32_t channels, - uint32_t applyActivation) { - causal_conv1d_batched_kernel<<>>( - input, output, weights, bias, batch, seqLen, channels, applyActivation); -} - -extern "C" void call_kernel_batched_bf16(uint32_t blockDim, void* stream, - uint8_t* input, uint8_t* output, - uint8_t* weights, uint8_t* bias, - uint32_t batch, uint32_t seqLen, - uint32_t channels, - uint32_t applyActivation) { - causal_conv1d_batched_bf16_kernel<<>>( - input, output, weights, bias, batch, seqLen, channels, applyActivation); -} diff --git a/tests/test_gdn_causal_conv1d.py b/tests/test_gdn_causal_conv1d.py new file mode 100644 index 00000000..42e72dc5 --- /dev/null +++ b/tests/test_gdn_causal_conv1d.py @@ -0,0 +1,87 @@ +# -------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# All rights reserved. +# See LICENSE in the root of the software repository: +# https://github.com/huawei-csl/pto-kernels/ +# for the full License text. +# -------------------------------------------------------------------------------- + +"""Correctness tests for the depthwise causal conv1d + bias + SiLU PTO kernel. + +Run: + pytest tests/test_gdn_causal_conv1d.py -q +""" + +import pytest +import torch +import torch.nn.functional as F + +from pto_kernels import pto_gdn_causal_conv1d, pto_gdn_causal_conv1d_batched + +# Filter width K compiled into the wheel (matches CAUSAL_CONV_K default). +K = 4 + +DTYPES = [torch.float16, torch.bfloat16] +# (batch, seq, dim): small general shapes + the GDN prefill regime +# (dim 2048 = H*D, 6144 = q+k+v; seq 128..512). +TEST_CASES = [ + (1, 16, 256), + (2, 31, 256), + (1, 128, 2048), + (8, 128, 2048), + (1, 256, 2048), + (8, 384, 2048), + (1, 512, 2048), + (8, 512, 2048), + (1, 128, 6144), + (8, 256, 6144), + (1, 384, 6144), + (8, 512, 6144), +] +# max abs error tolerance vs the fp32 reference (dtype rounding only). +TOL = {torch.float16: 6e-3, torch.bfloat16: 6e-2} + + +def gdn_causal_conv1d_ref(x, w, bias, activation): + """Depthwise causal conv1d (width K) + per-channel bias + optional SiLU. + + fp32 accumulate; x[:, <0] padded with zeros (no conv_states). + x: [B, L, W] w: [K, W] bias: [W] -> [B, L, W] (fp32) + """ + B, L, W = x.shape + pad = torch.zeros((B, K - 1, W), device=x.device, dtype=x.dtype) + xe = torch.cat([pad, x], dim=1).float() + wf = w.float() + acc = sum(xe[:, k : k + L] * wf[k] for k in range(K)) + bias.float() + return F.silu(acc) if activation else acc + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("activation", [True, False]) +@pytest.mark.parametrize("batch,seq,dim", TEST_CASES) +def test_batched_matches_reference(npu_device, dtype, activation, batch, seq, dim): + x = 2 * torch.rand(batch, seq, dim, device=npu_device, dtype=dtype) - 1 + w = torch.rand(K, dim, device=npu_device, dtype=torch.float32) - 0.5 + bias = torch.rand(dim, device=npu_device, dtype=torch.float32) - 0.5 + + y = pto_gdn_causal_conv1d_batched(x, w, bias, activation=activation) + torch.npu.synchronize() + + ref = gdn_causal_conv1d_ref(x, w, bias, activation) + err = (y.float() - ref).abs().max().item() + assert err <= TOL[dtype], f"max abs err {err:.3e} > tol {TOL[dtype]:.1e}" + + +@pytest.mark.parametrize("seq,dim", [(16, 256), (128, 2048), (512, 6144)]) +def test_single_fp16_entry_matches_reference(npu_device, seq, dim): + """The non-batched [L, W] fp16 entry (always applies SiLU).""" + x = 2 * torch.rand(seq, dim, device=npu_device, dtype=torch.float16) - 1 + w = torch.rand(K, dim, device=npu_device, dtype=torch.float32) - 0.5 + bias = torch.rand(dim, device=npu_device, dtype=torch.float32) - 0.5 + + y = pto_gdn_causal_conv1d(x, w, bias) + torch.npu.synchronize() + + ref = gdn_causal_conv1d_ref(x.unsqueeze(0), w, bias, activation=True)[0] + err = (y.float() - ref).abs().max().item() + assert err <= TOL[torch.float16], f"max abs err {err:.3e}" From 5bf05dde267e2afee2536ccad98962372aa2e405 Mon Sep 17 00:00:00 2001 From: anastasios Date: Sat, 27 Jun 2026 06:22:20 +0000 Subject: [PATCH 3/4] fix --- csrc/host/torch_gdn_causal_conv1d.h | 12 +-- csrc/kernel/kernel_gdn_causal_conv1d.cpp | 108 ++++++++++++----------- 2 files changed, 63 insertions(+), 57 deletions(-) diff --git a/csrc/host/torch_gdn_causal_conv1d.h b/csrc/host/torch_gdn_causal_conv1d.h index 00153d0c..664611e3 100644 --- a/csrc/host/torch_gdn_causal_conv1d.h +++ b/csrc/host/torch_gdn_causal_conv1d.h @@ -11,9 +11,9 @@ for the full License text. #include #include -#include "aclrtlaunch_gdn_causal_conv1d_batched_bf16_kernel.h" -#include "aclrtlaunch_gdn_causal_conv1d_batched_kernel.h" -#include "aclrtlaunch_gdn_causal_conv1d_kernel.h" +#include "aclrtlaunch_gdn_causal_conv1d_batched_bf16.h" +#include "aclrtlaunch_gdn_causal_conv1d_batched_fp16.h" +#include "aclrtlaunch_gdn_causal_conv1d_fp16.h" #include "utils.h" namespace pto_isa_ops { @@ -55,7 +55,7 @@ at::Tensor run_gdn_causal_conv1d(const at::Tensor& x, const at::Tensor& weights, at::Tensor output = at::empty_like(x); const uint32_t block_dim = GetNumVectorCores(); - EXEC_KERNEL_CMD(gdn_causal_conv1d_kernel, block_dim, x, output, weights, bias, + EXEC_KERNEL_CMD(gdn_causal_conv1d_fp16, block_dim, x, output, weights, bias, seqLen, channels); return output; } @@ -108,10 +108,10 @@ at::Tensor run_gdn_causal_conv1d_batched(const at::Tensor& x, const uint32_t block_dim = GetNumVectorCores(); if (x.scalar_type() == at::kHalf) { - EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_kernel, block_dim, x, output, + EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_fp16, block_dim, x, output, weights, bias, batch, seqLen, channels, applyActivation); } else if (x.scalar_type() == at::kBFloat16) { - EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_bf16_kernel, block_dim, x, output, + EXEC_KERNEL_CMD(gdn_causal_conv1d_batched_bf16, block_dim, x, output, weights, bias, batch, seqLen, channels, applyActivation); } return output; diff --git a/csrc/kernel/kernel_gdn_causal_conv1d.cpp b/csrc/kernel/kernel_gdn_causal_conv1d.cpp index ce5a47c5..4e8e2519 100644 --- a/csrc/kernel/kernel_gdn_causal_conv1d.cpp +++ b/csrc/kernel/kernel_gdn_causal_conv1d.cpp @@ -10,53 +10,68 @@ for the full License text. #include "kernel_utils.h" using namespace pto; +using namespace kernel_utils; + +// Filter width / per-tile channel width the entry points below are compiled at. +// Default to the K=4, MAX_W=3072 production configuration; the test suite +// recompiles this file at other widths via -DCAUSAL_CONV_K / +// -DCAUSAL_CONV_MAX_W (see test_gdn_causal_conv1d.py), so a plain build is +// byte-for-byte unaffected. +#ifndef CAUSAL_CONV_K +#define CAUSAL_CONV_K 4 +#endif +#ifndef CAUSAL_CONV_MAX_W +#define CAUSAL_CONV_MAX_W 3072 +#endif #define DIV_ROUNDUP(x, y) (((x) + (y) - 1) / (y)) #define ALIGN_UP(x, y) (DIV_ROUNDUP((x), (y)) * (y)) // =========================================================================== -// DEPTHWISE fused causal conv1d + bias + (optional) SiLU (per-channel, any K) +// DEPTHWISE fused causal conv1d + bias + (optional) SiLU (per-channel, any +// K) // -// y[b,i,c] = act( bias[c] + sum_{k=max(0,K-1-i)..K-1} W[k,c] * x[b, i-K+1+k, -// c] ), x[<0]=0 +// y[b,i,c] = act( bias[c] + sum_{k=max(0,K-1-i)..K-1} W[k,c] * x[b, +// i-K+1+k, c] ), x[<0]=0 // // Per-channel K-tap filter (Mamba/GDN short conv). x,y are [batch, seqLen, // channels] row-major; `channels` is the lane axis, seqLen the conv axis. -// Weights W[K,channels] + bias[channels] are fp32 GM tensors. fp16 OR bf16 I/O, -// fp32 accumulate. +// Weights W[K,channels] + bias[channels] are fp32 GM tensors. fp16 OR bf16 +// I/O, fp32 accumulate. // -// Filter width K and per-tile channel width MAX_W are compile-time constants -// chosen at the call site as template parameters (no preprocessor config), e.g. +// Filter width K and per-tile channel width MAX_W are compile-time +// constants chosen at the call site as template parameters (no preprocessor +// config), e.g. // constexpr uint32_t K = CAUSAL_CONV_K, MAX_W = CAUSAL_CONV_MAX_W; // csilu::runConvSiluBatched(...); // // 2-D-plus-batch work grid: workUnits = batch x sequenceChunkCount x // channelTileCount. Each work unit produces outputs [batchIndex] x // [outputRowStart,outputRowEnd) for channels -// [channelTileBase,channelTileBase+tileChannelCount), replaying K-1 causal halo -// rows to prime its accumulators. The grid fills all cores: batch supplies -// parallelism first, then channel tiles, then sequence chunks; the channel-tile -// width is widened to whole vector lanes for coalesced stores. (See -// processWorkUnit + runConvSiluBatched.) +// [channelTileBase,channelTileBase+tileChannelCount), replaying K-1 causal +// halo rows to prime its accumulators. The grid fills all cores: batch +// supplies parallelism first, then channel tiles, then sequence chunks; the +// channel-tile width is widened to whole vector lanes for coalesced stores. +// (See processWorkUnit + runConvSiluBatched.) // -// Generic in K (filter width) and MAX_W (tile width). accumRingSize = smallest -// power of two >= K is the accumulator ring size, so the K outputs in flight -// map to distinct ring slots via `idx & accumRingMask` (accumRingMask = -// accumRingSize - 1). -// UB layout (per lane, fp32 unless noted): K weights + bias(1) + accumRingSize -// accumulators + (K-1) partial products + input-as-fp32(1) = -// 2*K+accumRingSize+1 fp32 tiles; then the I/O region inputTile[0..1] + output0 -// + output1 = 4 I/O tiles (input load double-buffered). A static_assert keeps -// the total within UB_BYTES_PER_CORE for the chosen K / MAX_W / dtypes. NOTE: -// uses the PTO tile-op API (); the `csilu` namespace avoids a -// clash with pto::detail. +// Generic in K (filter width) and MAX_W (tile width). accumRingSize = +// smallest power of two >= K is the accumulator ring size, so the K outputs +// in flight map to distinct ring slots via `idx & accumRingMask` +// (accumRingMask = accumRingSize - 1). UB layout (per lane, fp32 unless +// noted): K weights + bias(1) + accumRingSize accumulators + (K-1) partial +// products + input-as-fp32(1) = 2*K+accumRingSize+1 fp32 tiles; then the +// I/O region inputTile[0..1] + output0 +// + output1 = 4 I/O tiles (input load double-buffered). A static_assert +// keeps the total within UB_BYTES_PER_CORE for the chosen K / MAX_W / +// dtypes. NOTE: uses the PTO tile-op API (); the `csilu` +// namespace avoids a clash with pto::detail. // =========================================================================== namespace csilu { // Unified Buffer available per AIV core (Ascend 910B2 = 192 KiB). The UB -// static_assert in processWorkUnit checks the chosen layout fits; raise it for -// a next-gen NPU with a larger UB. +// static_assert in processWorkUnit checks the chosen layout fits; raise it +// for a next-gen NPU with a larger UB. constexpr uint32_t UB_BYTES_PER_CORE = 192u * 1024u; // Smallest power of two >= value. @@ -163,7 +178,8 @@ AICORE inline void processWorkUnit( // double-buffered input: two load slots with independent handshakes. // EVENT_ID3 is reused here (the weight load above already consumed it). - // events are passed by mutable value to set_flag/wait_flag, so not constexpr. + // events are passed by mutable value to set_flag/wait_flag, so not + // constexpr. const event_t inputBufferEvent[2] = {EVENT_ID0, EVENT_ID3}; set_flag(PIPE_V, PIPE_MTE2, inputBufferEvent[0]); // input slot 0 initially free @@ -193,8 +209,8 @@ AICORE inline void processWorkUnit( TASSIGN(inputTileIo, ubInputOffset[bufferIndex]); TASSIGN(inputTileFp32, ubInputFp32Offset); - // (1) consume current row from buffer bufferIndex (loaded by the prologue / - // previous prefetch). + // (1) consume current row from buffer bufferIndex (loaded by the prologue + // / previous prefetch). wait_flag(PIPE_MTE2, PIPE_V, inputBufferEvent[bufferIndex]); TCVT(inputTileFp32, inputTileIo, pto::RoundMode::CAST_NONE); set_flag(PIPE_V, PIPE_MTE2, @@ -218,14 +234,16 @@ AICORE inline void processWorkUnit( pipe_barrier(PIPE_V); // scatter: this input row contributes to K output rows; form each - // weight*input product (only for outputs in [outputRowStart,outputRowEnd)). + // weight*input product (only for outputs in + // [outputRowStart,outputRowEnd)). for (uint32_t tapIndex = 0; tapIndex < K; ++tapIndex) { const uint32_t outputRow = inputRow + (K - 1) - tapIndex; if (outputRow < outputRowStart || outputRow >= outputRowEnd) continue; AccumTile weightTile(tileChannelCount); TASSIGN(weightTile, tapIndex * accumTileBytes); if (inputRow == 0 || tapIndex == 0) { - // first contribution to this output's accumulator slot: initialise it. + // first contribution to this output's accumulator slot: initialise + // it. AccumTile accumTile(tileChannelCount); TASSIGN(accumTile, ubAccumRingBase + (outputRow & accumRingMask) * accumTileBytes); @@ -309,9 +327,9 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, if (seqLen == 0 || batch == 0 || channels == 0) return; // ---- grid tuning knobs ---- - // Minimum rows per sequence-chunk. Smaller chunks expose more parallelism but - // each replays K-1 causal halo rows, so very small chunks waste compute; 32 - // is the balance point across the GDN sequence lengths. + // Minimum rows per sequence-chunk. Smaller chunks expose more parallelism + // but each replays K-1 causal halo rows, so very small chunks waste + // compute; 32 is the balance point across the GDN sequence lengths. constexpr uint32_t minSeqChunkLen = 32u; // Channels per aligned vector lane = 256 B / element size (128 for // fp16/bf16). Channel tiles are sized in whole lanes for coalesced @@ -327,8 +345,8 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, uint32_t maxSequenceChunks = DIV_ROUNDUP(seqLen, minSeqChunkLen); if (maxSequenceChunks < 1) maxSequenceChunks = 1; - // channel tiles: enough that each tile fits MAX_W, and enough (with the chunk - // count) to fill the per-sequence work-unit budget. + // channel tiles: enough that each tile fits MAX_W, and enough (with the + // chunk count) to fill the per-sequence work-unit budget. const uint32_t channelTilesForUbLimit = DIV_ROUNDUP(channels, MAX_W); const uint32_t channelTilesToFillCores = DIV_ROUNDUP(workUnitsPerSequence, maxSequenceChunks); @@ -384,18 +402,6 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, } // namespace csilu -// Filter width / per-tile channel width the entry points below are compiled at. -// Default to the K=4, MAX_W=3072 production configuration; the test suite -// recompiles this file at other widths via -DCAUSAL_CONV_K / -// -DCAUSAL_CONV_MAX_W (see test_gdn_causal_conv1d.py), so a plain build is -// byte-for-byte unaffected. -#ifndef CAUSAL_CONV_K -#define CAUSAL_CONV_K 4 -#endif -#ifndef CAUSAL_CONV_MAX_W -#define CAUSAL_CONV_MAX_W 3072 -#endif - /** * @brief Single-sequence causal depthwise conv1d with optional SiLU activation * (fp16 I/O). @@ -423,7 +429,7 @@ AICORE void runConvSiluBatched(__gm__ IoElemType* input, * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On * other targets all parameters are ignored and the kernel is a no-op. */ -extern "C" __global__ AICORE void gdn_causal_conv1d_kernel( +extern "C" __global__ AICORE void gdn_causal_conv1d_fp16( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t seqLen, uint32_t channels) { #if defined(__DAV_VEC__) @@ -467,7 +473,7 @@ extern "C" __global__ AICORE void gdn_causal_conv1d_kernel( * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On * other targets all parameters are ignored and the kernel is a no-op. */ -extern "C" __global__ AICORE void gdn_causal_conv1d_batched_kernel( +extern "C" __global__ AICORE void gdn_causal_conv1d_batched_fp16( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, uint32_t applyActivation) { @@ -492,7 +498,7 @@ extern "C" __global__ AICORE void gdn_causal_conv1d_batched_kernel( * @brief Batched causal depthwise conv1d with optional SiLU activation * (bfloat16 I/O). * - * Identical in semantics to @c gdn_causal_conv1d_batched_kernel but uses + * Identical in semantics to @c gdn_causal_conv1d_batched_fp16 but uses * bfloat16 for input and output tensors instead of fp16. Weights and bias * remain fp32. Filter width and maximum sequence length are fixed at compile * time via @c CAUSAL_CONV_K and @@ -516,7 +522,7 @@ extern "C" __global__ AICORE void gdn_causal_conv1d_batched_kernel( * @note Only compiled and executed on DAV vector cores (@c __DAV_VEC__). On * other targets all parameters are ignored and the kernel is a no-op. */ -extern "C" __global__ AICORE void gdn_causal_conv1d_batched_bf16_kernel( +extern "C" __global__ AICORE void gdn_causal_conv1d_batched_bf16( __gm__ uint8_t* input, __gm__ uint8_t* output, __gm__ uint8_t* weights, __gm__ uint8_t* bias, uint32_t batch, uint32_t seqLen, uint32_t channels, uint32_t applyActivation) { From 688da6f51cf62ae23527dfbccec48feda31524e2 Mon Sep 17 00:00:00 2001 From: anastasios Date: Mon, 29 Jun 2026 08:18:32 +0000 Subject: [PATCH 4/4] fix kernel for A5 --- csrc/kernel/kernel_gdn_causal_conv1d.cpp | 16 ++++++++-------- csrc/kernel/kernel_utils.h | 10 ++++++++++ 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/csrc/kernel/kernel_gdn_causal_conv1d.cpp b/csrc/kernel/kernel_gdn_causal_conv1d.cpp index 4e8e2519..f1d9369f 100644 --- a/csrc/kernel/kernel_gdn_causal_conv1d.cpp +++ b/csrc/kernel/kernel_gdn_causal_conv1d.cpp @@ -93,11 +93,11 @@ template AICORE inline void applySiluToTile(TileT& dst, TileT& src, TileT& scratch) { using ElemType = typename TileT::DType; TMULS(scratch, src, (ElemType)-1); - pipe_barrier(PIPE_V); + PipeBarrierVec(); TEXP(scratch, scratch); - pipe_barrier(PIPE_V); + PipeBarrierVec(); TADDS(scratch, scratch, (ElemType)1); - pipe_barrier(PIPE_V); + PipeBarrierVec(); TDIV(dst, src, scratch); } @@ -231,7 +231,7 @@ AICORE inline void processWorkUnit( set_flag(PIPE_MTE2, PIPE_V, inputBufferEvent[nextBufferIndex]); } - pipe_barrier(PIPE_V); + PipeBarrierVec(); // scatter: this input row contributes to K output rows; form each // weight*input product (only for outputs in @@ -254,7 +254,7 @@ AICORE inline void processWorkUnit( TMUL(productTile, inputTileFp32, weightTile); } } - pipe_barrier(PIPE_V); + PipeBarrierVec(); if (inputRow != 0) { for (uint32_t tapIndex = 1; tapIndex < K; ++tapIndex) { const uint32_t outputRow = inputRow + (K - 1) - tapIndex; @@ -267,7 +267,7 @@ AICORE inline void processWorkUnit( TADD(accumTile, accumTile, productTile); } } - pipe_barrier(PIPE_V); + PipeBarrierVec(); if (inputRow < outputRowStart) continue; // halo row: primed accumulators only @@ -285,10 +285,10 @@ AICORE inline void processWorkUnit( TASSIGN(outputTile, ubOutputOffset[outputBufferIndex]); TADD(accumTile, accumTile, biasTile); - pipe_barrier(PIPE_V); + PipeBarrierVec(); if (applyActivation) { applySiluToTile(accumTile, accumTile, siluScratchTile); - pipe_barrier(PIPE_V); + PipeBarrierVec(); } wait_flag(PIPE_MTE3, PIPE_V, outputBufferEvent); TCVT(outputTile, accumTile, pto::RoundMode::CAST_NONE); diff --git a/csrc/kernel/kernel_utils.h b/csrc/kernel/kernel_utils.h index 4f60d2ec..c5bee836 100644 --- a/csrc/kernel/kernel_utils.h +++ b/csrc/kernel/kernel_utils.h @@ -165,4 +165,14 @@ constexpr pto::BLayout GetOuterLayout(bool is_left) { #endif } +/** + * @brief Pipe in-core barrier for vector core that is safe to use with A2A3 and + * A5. On A5, PipeBarrierVec() is a noop. + */ +AICORE inline void PipeBarrierVec() { +#if __CCE_AICORE__ == 220 + pipe_barrier(PIPE_V); +#endif +} + } // namespace kernel_utils