More validation + fixup tests

Author: christiangnrd

src/fft.jl

@@ -78,8 +78,11 @@ mutable struct MtlFFTPlan{T <: FFTNumber, S <: FFTNumber, backward, inplace, N,
 
     function MtlFFTPlan{T, S, backward, inplace, N, R}(input_size::NTuple{N, Int}, output_size::NTuple{N, Int}, region::NTuple{R, Int}) where {T <: FFTNumber, S <: FFTNumber, backward, inplace, N, R}
         # Validate region
-        for r in region
-            1 <= r <= N || throw(ArgumentError("Invalid FFT dimension $r for array with $N dimensions"))
+        if any(diff(collect(region)) .< 1)
+            throw(ArgumentError("region must be an increasing sequence"))
+        end
+        if any(region .< 1 .|| region .> N)
+            throw(ArgumentError("region can only refer to valid dimensions"))
         end
         backward isa Bool || throw(ArgumentError("FFT backward argument must be a Bool"))
         inplace isa Bool || throw(ArgumentError("FFT inplace argument must be a Bool"))
@@ -277,14 +280,15 @@ function LinearAlgebra.mul!(y::MtlArray{T, N}, p::MtlFFTPlan{T, S, backward, inp
 end
 
 function Base.:(*)(p::MtlFFTPlan{T, S, backward, true}, x::MtlArray{S}) where {T, S, backward}
-    # assert_applicable(p, x)
-    LinearAlgebra.mul!(x, p, x)
+    assert_applicable(p, x)
+
+    unsafe_execute!(p, x, x)
     return x
 end
 function Base.:(*)(p::MtlFFTPlan{T, S, backward, false}, x::MtlArray{S}) where {T, S, backward}
-    # assert_applicable(p, x)
+    assert_applicable(p, x)
 
     y = MtlArray{T}(undef, p.output_size)
-    LinearAlgebra.mul!(y, p, x)
+    unsafe_execute!(p, x, y)
     return y
 end

test/fft.jl

@@ -159,8 +159,8 @@ if MPS.is_supported(device())
 
         p = plan_fft(d_X, region)
         d_Y = p * d_X
-        # d_X2 = reshape(d_X, (size(d_X)..., 1))
-        # @test_throws ArgumentError p * d_X2
+        d_X2 = reshape(d_X, (size(d_X)..., 1))
+        @test_throws ArgumentError p * d_X2
 
         Y = Array(d_Y)
         @test isapprox(Y, fftw_X, rtol = rtol(T), atol = atol(T))
@@ -238,19 +238,17 @@ if MPS.is_supported(device())
                     test_complex_batched(X, region)
                 end
 
-                # X = rand(T, dims)
-                # @test_throws ArgumentError test_complex_batched(X, (3, 1))
+                X = rand(T, dims)
+                @test_throws ArgumentError test_complex_batched(X, (3, 1))
             end
             @testset "Batch 2D (in 4D)" begin
                 dims = (N1, N2, N3, N4)
                 for region in [(1, 2), (1, 4), (3, 4), (1, 3), (2, 3), (2,), (3,)]
                     X = rand(T, dims)
                     test_complex_batched(X, region)
                 end
-                # for region in [(2, 4)]
-                #     X = rand(T, dims)
-                #     @test_throws ArgumentError test_complex_batched(X, region)
-                # end
+                X = rand(T, dims)
+                test_complex_batched(X, (2, 4))
             end
         end
     end
@@ -338,8 +336,8 @@ if MPS.is_supported(device())
                     test_real_batched(X, region)
                 end
 
-            #     X = rand(T, dims)
-            #     @test_throws ArgumentError test_real_batched(X, (3, 1))
+                X = rand(T, dims)
+                @test_throws ArgumentError test_real_batched(X, (3, 1))
             end
 
             @testset "Batch 2D (in 4D)" begin
@@ -348,10 +346,8 @@ if MPS.is_supported(device())
                     X = rand(T, dims)
                     test_real_batched(X, region)
                 end
-            #     for region in [(2,4)]
-            #         X = rand(T, dims)
-            #         @test_throws ArgumentError test_real_batched(X, region)
-            #     end
+                X = rand(T, dims)
+                test_real_batched(X, (2, 4))
             end
 
             @testset "3D" begin