Skip to content

NKI ComplexConv1d kernel #49

Description

@scttfrdmn

Today ComplexConv1d.forward always falls back to 2× torch.nn.Conv1d calls even when set_backend("nki"). No NKI dispatch path exists (only ComplexLinear has one).

Options:

  1. im2col-style kernel that unrolls the conv into a GEMM, reuses _complex_gemm_kernel
  2. Native conv kernel with direct windowed memory access

Motivated by speech enhancement workloads (cIRM) that use ComplexConv1d in the front-end.

Metadata

Metadata

Assignees

No one assigned

    Labels

    nkiNKI kernel development (Trainium hardware)nnComplex neural network layers

    Type

    No type

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions