kernel source

27813b0 verified 17 days ago

5.13 kB

	// Polyphase rational resampling.
	//
	// Converting by L/M is, in principle, zero-stuff by L, low-pass, keep every
	// Mth sample. Done literally that computes L-1 zeros for every real sample and
	// then throws away M-1 of every M results. The polyphase identity removes both:
	// output n depends only on filter phase (n*M) mod L and a short run of input
	// around floor(n*M/L).
	//
	// torchaudio expresses this as a conv1d with `new_freq` output channels and
	// stride `orig_freq`. That is the same decomposition, but the per-phase filters
	// are stored padded to a common width, so every output sample multiplies through
	// a kernel row whose support is a small fraction of its length. At 44.1k -> 16k
	// with the default filter width the row is 475 taps and the support is about 34,
	// and the padding is multiplied and summed like everything else.
	//
	// Here each phase keeps only its own support, as an offset and a run of taps.

	#include <torch/all.h>

	#include <algorithm>
	#include <cstdint>
	#include <cstring>
	#include <vector>

	namespace {

	#if defined(__GNUC__) \|\| defined(__clang__)
	#if defined(__x86_64__)
	#define RS_VL 8
	#else
	#define RS_VL 4
	#endif
	typedef float vf __attribute__((vector_size(RS_VL * sizeof(float))));
	#else
	#define RS_VL 1
	typedef float vf;
	#endif

	#if defined(__x86_64__) && defined(__GNUC__) && !defined(__clang__)
	#define RS_MV __attribute__((target_clones("avx2", "default"), flatten))
	#else
	#define RS_MV
	#endif

	inline vf vzero() {
	#if RS_VL == 1
	return 0.0f;
	#else
	vf v;
	float* p = reinterpret_cast<float*>(&v);
	for (int i = 0; i < RS_VL; i++) p[i] = 0.0f;
	return v;
	#endif
	}

	inline vf vloadu(const float* p) {
	vf v;
	std::memcpy(&v, p, sizeof(vf));
	return v;
	}

	// Tree reduction rather than a running scalar sum. Phases here carry as few as
	// a dozen taps, so the dot product is two or three vector operations and the
	// reduction that follows it is a large share of the per-output cost; a chain of
	// RS_VL dependent adds is the wrong shape for that.
	inline float hsum(const vf& v) {
	#if RS_VL == 1
	return v;
	#else
	float t[RS_VL];
	std::memcpy(t, &v, sizeof(vf));
	for (int s = RS_VL / 2; s > 0; s >>= 1)
	for (int i = 0; i < s; i++) t[i] += t[i + s];
	return t[0];
	#endif
	}

	// One output block: the L outputs sharing input block `a`. Each reads a
	// contiguous run of input, so the tap loop is unit stride in both operands.
	RS_MV void resample_block(float* out, const float* x, int64_t T, int64_t a,
	int64_t L, int64_t M, int64_t n_out,
	const int32_t* start, const int32_t* len,
	const int64_t* off, const float* taps) {
	const int64_t abase = a * M;
	for (int64_t r = 0; r < L; r++) {
	const int64_t n = a * L + r;
	if (n >= n_out) return;
	const int64_t base = abase + start[r];
	const int64_t K = len[r];
	const float* h = taps + off[r];

	// Clip the tap run to where the input index is in range rather than
	// testing every tap; the interior blocks take the fast path untouched.
	int64_t m0 = 0, m1 = K;
	if (base < 0) m0 = -base;
	if (base + K > T) m1 = T - base;
	if (m0 >= m1) { out[n] = 0.0f; continue; }

	const float* xp = x + base;
	vf acc = vzero();
	int64_t m = m0;
	for (; m + RS_VL <= m1; m += RS_VL)
	acc = acc + vloadu(h + m) * vloadu(xp + m);
	float s = hsum(acc);
	for (; m < m1; m++) s += h[m] * xp[m];
	out[n] = s;
	}
	}

	} // namespace

	// out: [B, n_out] f32
	// x: [B, T] f32 contiguous
	// taps: concatenated per-phase filter supports, f32
	// start: [L] int32, input offset of each phase's first tap relative to a*M
	// len: [L] int32, tap count per phase
	// off: [L] int64, index of each phase's first tap within `taps`
	void rp_resample(torch::Tensor& out, torch::Tensor const& x,
	torch::Tensor const& taps, torch::Tensor const& start,
	torch::Tensor const& len, torch::Tensor const& off,
	int64_t L, int64_t M) {
	TORCH_CHECK(x.dim() == 2 && x.is_contiguous() &&
	x.scalar_type() == torch::kFloat32,
	"x must be contiguous f32 [B, T]");
	TORCH_CHECK(out.dim() == 2 && out.is_contiguous(), "out must be [B, n_out]");
	TORCH_CHECK(L >= 1 && M >= 1, "L and M must be positive");
	TORCH_CHECK(start.numel() == L && len.numel() == L && off.numel() == L,
	"phase tables must have L entries");

	const int64_t B = x.size(0), T = x.size(1), n_out = out.size(1);
	const float* xp = x.data_ptr<float>();
	float* op = out.data_ptr<float>();
	const float* tp = taps.data_ptr<float>();
	const int32_t* sp = start.data_ptr<int32_t>();
	const int32_t* lp = len.data_ptr<int32_t>();
	const int64_t* fp = off.data_ptr<int64_t>();

	const int64_t n_blocks = (n_out + L - 1) / L;

	at::parallel_for(0, B * n_blocks, 1, [&](int64_t begin, int64_t end) {
	for (int64_t i = begin; i < end; i++) {
	const int64_t b = i / n_blocks;
	const int64_t a = i - b * n_blocks;
	resample_block(op + b * n_out, xp + b * T, T, a, L, M, n_out,
	sp, lp, fp, tp);
	}
	});
	}