Using Python Taichi to render audio on the GPU and output via Sound Device

taichi_audio.py

from iipyper import Audio, run
import sounddevice as sd
import taichi as ti
from math import pi
import sys
@ti.kernel
def generate_data(outdata: ti.ext_arr(), 
                  frames: ti.i32, 
                  start_idx: ti.template(), 
                  fs: ti.f32):
    for I in ti.grouped(ti.ndrange(frames)):
        i = I[0]
        t = (start_idx[None] + i) / fs
        s = ti.sin(2 * pi * 440 * t)
        start_idx[None] += 1
        outdata[i,0] = s
def main():
    ti.init(arch=ti.vulkan)
    device = sd.default.device
    fs = sd.query_devices(sd.default.device, 'output')['default_samplerate']
    buffer_size = 128
    start_idx = ti.field(ti.i32, ())
    def callback(indata, outdata, frames, time, status):
        if status: print(status, file=sys.stderr)
        generate_data(outdata, frames, start_idx, fs)
    Audio(device=device, channels=1, blocksize=buffer_size,
        samplerate=fs, callback=callback)
if __name__=='__main__':
    run(main)