Add convolution support (DSP.conv / xcorr) via the DSP.jl interface

Project.toml

@@ -33,9 +33,11 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [weakdeps]
+DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
+MetalDSPExt = "DSP"
 SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -44,6 +46,7 @@ Adapt = "4.5"
 BFloat16s = "0.5, 0.6"
 CEnum = "0.4, 0.5"
 CodecBzip2 = "0.8.5"
+DSP = "0.7, 0.8"
 ExprTools = "0.1"
 GPUArrays = "11.5"
 GPUCompiler = "1.13.2"

ext/MetalDSPExt.jl

@@ -0,0 +1,60 @@
+module MetalDSPExt
+
+# GPU-accelerated linear convolution and cross-correlation for `MtlArray`s.
+#
+# Dispatches `DSP.conv` / `DSP.xcorr` to Metal's FFT-based convolution engine
+# (`Metal.MPSGraphs`), so `using DSP, Metal; conv(a, b)` runs on the GPU instead
+# of falling back to scalar CPU indexing. This mirrors how the FFT support
+# extends `AbstractFFTs` rather than introducing a parallel `Metal.conv`.
+
+using Metal
+import DSP
+
+# Element types the MPSGraph FFT/convolution engine supports.
+const MtlConvNumber = Union{Float32, Float16, ComplexF32, ComplexF16}
+
+# Clear error for unsupported element types (e.g. Float64). The MtlArray methods
+# below intercept DSP.conv/xcorr before DSP's generic CPU path, which would hit
+# disallowed scalar indexing on the device.
+@noinline function _conv_unsupported(u, v)
+    throw(
+        ArgumentError(
+            "Metal convolution supports Float32/Float16 and their Complex types; " *
+                "got $(eltype(u)) and $(eltype(v)). Convert first, e.g. with `Float32.(x)`."
+        )
+    )
+end
+
+"""
+    DSP.conv(u::MtlArray, v::MtlArray; algorithm = :auto)
+
+Full linear convolution of two `MtlArray`s on the GPU, computed via the FFT
+convolution theorem. Convolves over all dimensions, matching `DSP.conv`
+semantics. The `algorithm` keyword (`:auto`/`:fft`) is accepted for
+compatibility; the FFT path is always used.
+"""
+function DSP.conv(
+        u::MtlArray{T, N}, v::MtlArray{T, N}; algorithm::Symbol = :auto
+    ) where {T <: Number, N}
+    T <: MtlConvNumber || _conv_unsupported(u, v)
+    # `algorithm` accepted for DSP.conv compatibility; the FFT engine is always used.
+    return Metal.MPSGraphs.conv(u, v; dims = ntuple(identity, N), mode = :full)
+end
+
+"""
+    DSP.xcorr(u::MtlVector, v::MtlVector; padmode = :none, scaling = :none)
+
+Cross-correlation of two GPU vectors, conjugating `v` (the DSP/MATLAB
+convention). Only `padmode = :none` (the full correlation) and `scaling = :none`
+are supported.
+"""
+function DSP.xcorr(
+        u::MtlVector{T}, v::MtlVector{T}; padmode::Symbol = :none, scaling::Symbol = :none
+    ) where {T <: Number}
+    T <: MtlConvNumber || _conv_unsupported(u, v)
+    padmode === :none || throw(ArgumentError("MetalDSPExt supports only padmode = :none"))
+    scaling === :none || throw(ArgumentError("MetalDSPExt supports only scaling = :none"))
+    return Metal.MPSGraphs.xcorr(u, v; mode = :full)
+end
+
+end # module

lib/mpsgraphs/MPSGraphs.jl

@@ -49,5 +49,6 @@ include("random.jl")
 
 include("matmul.jl")
 include("fft.jl")
+include("convolution.jl")
 
 end