chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,214 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks, unused-argument
|
||||
"""STFT operator"""
|
||||
|
||||
from math import pi
|
||||
|
||||
from tvm import te, tirx
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
|
||||
|
||||
def stft(
|
||||
data,
|
||||
n_fft,
|
||||
hop_length,
|
||||
win_length,
|
||||
window,
|
||||
normalized,
|
||||
onesided,
|
||||
output_shape,
|
||||
):
|
||||
"""
|
||||
The STFT computes the Fourier transform of short overlapping windows of the input.
|
||||
This gives frequency components of the signal as they change over time.
|
||||
Parameters
|
||||
----------
|
||||
data : te.Tensor
|
||||
Either a 1-D tensor or a 2-D batch tensor.
|
||||
n_fft : int
|
||||
The size of Fourier transform
|
||||
hop_length : int
|
||||
The distance between neighboring sliding window frames
|
||||
win_length : int
|
||||
The size of window frame and STFT filter
|
||||
window : te.Tensor
|
||||
A 1-D tensor window frame
|
||||
normalized : bool
|
||||
Whether to return the normalized STFT results
|
||||
onesided : bool
|
||||
Whether to return onesided result or fill with conjugate symmetry
|
||||
Returns
|
||||
-------
|
||||
output : te.Tensor
|
||||
Tensor containing the STFT result
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
data = [1, 2, 3, 4, 5, 6]
|
||||
window = [4, 3, 2]
|
||||
[n_fft, hop_length, win_length, normalized, onesided] = [3, 3, 3, False, True]
|
||||
topi.stft(data, n_fft, hop_length, win_length, window, normalized, onesided)
|
||||
-> [[[15.0000, 0.0000], [34.0000, 0.0000]], [[ 4.5000, 0.8660], [ 1.0000, -1.7321]]]
|
||||
"""
|
||||
|
||||
def gen_ir(
|
||||
data_ptr,
|
||||
n_fft,
|
||||
hop_length,
|
||||
win_length,
|
||||
window_ptr,
|
||||
normalized,
|
||||
onesided,
|
||||
output_ptr,
|
||||
loop_kind,
|
||||
):
|
||||
col_loop = T.vectorized if loop_kind == "vectorize" else T.serial
|
||||
|
||||
with IRBuilder() as ib:
|
||||
data = T.buffer_proxy(data_ptr)
|
||||
window = T.buffer_proxy(window_ptr)
|
||||
output = T.buffer_proxy(output_ptr)
|
||||
# https://librosa.org/doc/0.7.2/_modules/librosa/core/spectrum.html#stft
|
||||
with T.parallel(0, output_ptr.shape[0] * output_ptr.shape[1]) as batch_row:
|
||||
with col_loop(0, output_ptr.shape[2]) as col:
|
||||
batch = tirx.floordiv(batch_row, output_ptr.shape[1])
|
||||
row = tirx.floormod(batch_row, output_ptr.shape[1])
|
||||
output[batch, row, col, 0] = tirx.Cast(data_ptr.dtype, 0)
|
||||
output[batch, row, col, 1] = tirx.Cast(data_ptr.dtype, 0)
|
||||
with T.serial(0, win_length) as wlen:
|
||||
output[batch, row, col, 0] += (
|
||||
window[wlen]
|
||||
* data[batch, col * hop_length + wlen]
|
||||
* tirx.cos(2 * pi * row * wlen / win_length)
|
||||
)
|
||||
output[batch, row, col, 1] -= (
|
||||
window[wlen]
|
||||
* data[batch, col * hop_length + wlen]
|
||||
* tirx.sin(2 * pi * row * wlen / win_length)
|
||||
)
|
||||
with T.If(normalized):
|
||||
with T.Then():
|
||||
output[batch, row, col, 0] /= tirx.sqrt(tirx.const(n_fft, "float32"))
|
||||
output[batch, row, col, 1] /= tirx.sqrt(tirx.const(n_fft, "float32"))
|
||||
|
||||
return ib.get()
|
||||
|
||||
output_buf = tirx.decl_buffer(output_shape, data.dtype, "output_buf", layout=None)
|
||||
loop_kind = "vectorize"
|
||||
if isinstance(output_shape[2], tirx.expr.Var): # any_dim
|
||||
loop_kind = "serial"
|
||||
|
||||
return te.extern(
|
||||
output_shape,
|
||||
[data, window],
|
||||
lambda ins, outs: gen_ir(
|
||||
ins[0], n_fft, hop_length, win_length, ins[1], normalized, onesided, outs[0], loop_kind
|
||||
),
|
||||
dtype=[data.dtype],
|
||||
out_buffers=[output_buf],
|
||||
name="stft_cpu",
|
||||
tag="stft_cpu",
|
||||
)
|
||||
|
||||
|
||||
def dft(
|
||||
re_data: te.Tensor,
|
||||
im_data: te.Tensor,
|
||||
inverse: tirx.IntImm,
|
||||
):
|
||||
"""
|
||||
Computes the discrete Fourier transform of input (calculation along the last axis).
|
||||
This gives frequency components of the signal as they change over time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
re_data : te.Tensor
|
||||
N-D tensor, real part of the input signal.
|
||||
|
||||
im_data : te.Tensor
|
||||
N-D tensor, imaginary part of the input signal.
|
||||
If the signal is real, then the values of this tensor are zeros.
|
||||
|
||||
inverse : bool
|
||||
Whether to perform the inverse discrete fourier transform.
|
||||
|
||||
Returns
|
||||
-------
|
||||
re_output : te.Tensor
|
||||
The Fourier Transform of the input (Real part).
|
||||
im_output : te.Tensor
|
||||
The Fourier Transform of the input (Imaginary part).
|
||||
"""
|
||||
|
||||
def gen_ir(
|
||||
re_data_buf,
|
||||
im_data_buf,
|
||||
re_output_buf,
|
||||
im_output_buf,
|
||||
):
|
||||
with IRBuilder() as ib:
|
||||
re_data_ptr = T.buffer_proxy(re_data_buf)
|
||||
im_data_ptr = T.buffer_proxy(im_data_buf)
|
||||
re_output_ptr = T.buffer_proxy(re_output_buf)
|
||||
im_output_ptr = T.buffer_proxy(im_output_buf)
|
||||
|
||||
shape = re_data.shape
|
||||
n_fft = shape[len(shape) - 1]
|
||||
base_range = 1
|
||||
for i in range(len(shape) - 1):
|
||||
base_range *= shape[i]
|
||||
|
||||
sign = -1 if inverse else 1
|
||||
factor = 1.0 / n_fft if inverse else 1.0
|
||||
|
||||
with T.parallel(0, base_range) as i:
|
||||
base_idx = i * n_fft
|
||||
with T.serial(0, n_fft) as n:
|
||||
n_idx = base_idx + n
|
||||
re_output_ptr[n_idx] = tirx.Cast(re_output_ptr.dtype, 0)
|
||||
im_output_ptr[n_idx] = tirx.Cast(im_output_ptr.dtype, 0)
|
||||
_w = sign * -2 * pi * n / n_fft
|
||||
with T.serial(0, n_fft) as k:
|
||||
k_idx = base_idx + k
|
||||
w = _w * k
|
||||
cos_w = tirx.Cast(re_output_ptr.dtype, tirx.cos(w))
|
||||
sin_w = tirx.Cast(re_output_ptr.dtype, tirx.sin(w))
|
||||
re_output_ptr[n_idx] += (
|
||||
re_data_ptr[k_idx] * cos_w - im_data_ptr[k_idx] * sin_w
|
||||
)
|
||||
im_output_ptr[n_idx] += (
|
||||
re_data_ptr[k_idx] * sin_w + im_data_ptr[k_idx] * cos_w
|
||||
)
|
||||
|
||||
re_output_ptr[n_idx] *= tirx.Cast(re_output_ptr.dtype, factor)
|
||||
im_output_ptr[n_idx] *= tirx.Cast(im_output_ptr.dtype, factor)
|
||||
|
||||
return ib.get()
|
||||
|
||||
output_shape = [re_data.shape] * 2
|
||||
|
||||
return te.extern(
|
||||
shape=output_shape,
|
||||
inputs=[re_data, im_data],
|
||||
fcompute=lambda ins, outs: gen_ir(ins[0], ins[1], outs[0], outs[1]),
|
||||
dtype=[re_data.dtype, im_data.dtype],
|
||||
name="dft_cpu",
|
||||
tag="dft_cpu",
|
||||
)
|
||||
Reference in New Issue
Block a user