fix(rocfft): include attach_halfN in the 2D twiddle table cache key

projects/rocfft/CHANGELOG.md

@@ -35,6 +35,12 @@ Documentation for rocFFT is available at
 ### Resolved issues
 
 * Addressed a cache-reuse issue with RCCL communicators by giving each communicator its own set of streams.
+* Fixed multi-dimensional real transforms returning wrong results when a complex
+  transform that shares the same internal 2D kernel had been planned earlier in the
+  process and its plan was still alive.  A real transform's kernel needs extra
+  twiddle values for its fused real pre/post-processing step, but the 2D twiddle
+  table cache key did not distinguish the two, so the real transform could be handed
+  the complex transform's shorter table.
 * Fixed `rocfft_plan_create` hanging when given a zero FFT length, zero batch, or zero dimensions; these
   now return `rocfft_status_invalid_dimensions` or `rocfft_status_invalid_arg_value`.
 * Fixed `rocfft_execution_info_set_stream` to derive the device from the stream itself instead of assuming the current device.

projects/rocfft/clients/tests/unit_test.cpp

@@ -473,6 +473,213 @@ TEST(rocfft_UnitTest, setup_cleanup_counter)
     ROCFFT_CATCH_TEST_EXCEPTIONS
 }
 
+// check that a multi-dimensional real inverse transform gets its own twiddle
+// table when a plain complex kernel of the same size already has one cached.
+// The real transform runs on a complex kernel with fused pre-processing, which
+// needs a half-N section appended that the complex kernel does not use.
+TEST(rocfft_UnitTest, twiddle_cache_2D_real_after_complex)
+{
+    PROB_SKIP_UNITTEST();
+    try
+    {
+        LocalCleanup restore([]() {
+            rocfft_cleanup();
+            rocfft_setup();
+        });
+
+        // a real length of 8 on one dimension and 4 on the other decomposes to
+        // a 4x4 complex kernel whichever dimension is the fastest, so it
+        // collides with the complex transform below either way
+        const std::vector<size_t> real_lengths{8, 4};
+        const std::vector<size_t> complex_lengths{4, 4};
+
+        const size_t input_elems  = (real_lengths[0] / 2 + 1) * real_lengths[1];
+        const size_t output_elems = real_lengths[0] * real_lengths[1];
+
+        // the input need not be Hermitian-symmetric - the property under test
+        // is only that the same input gives the same output
+        std::vector<rocfft_complex<float>> input(input_elems);
+        for(size_t i = 0; i < input.size(); ++i)
+            input[i] = {static_cast<float>(i % 7) - 3.0f, static_cast<float>(i % 5) - 2.0f};
+
+        auto real_inverse = [&]() {
+            const size_t                    ibytes = input.size() * sizeof(input[0]);
+            const size_t                    obytes = output_elems * sizeof(float);
+            gpubuf_t<rocfft_complex<float>> idev;
+            gpubuf_t<float>                 odev;
+            EXPECT_EQ(idev.alloc(ibytes), hipSuccess);
+            EXPECT_EQ(odev.alloc(obytes), hipSuccess);
+            EXPECT_EQ(hipMemcpy(idev.data(), input.data(), ibytes, hipMemcpyHostToDevice),
+                      hipSuccess);
+
+            rocfft_plan plan = nullptr;
+            EXPECT_EQ(rocfft_plan_create(&plan,
+                                         rocfft_placement_notinplace,
+                                         rocfft_transform_type_real_inverse,
+                                         rocfft_precision_single,
+                                         real_lengths.size(),
+                                         real_lengths.data(),
+                                         1,
+                                         nullptr),
+                      rocfft_status_success);
+            void* ibuf = idev.data();
+            void* obuf = odev.data();
+            EXPECT_EQ(rocfft_execute(plan, &ibuf, &obuf, nullptr), rocfft_status_success);
+            EXPECT_EQ(hipDeviceSynchronize(), hipSuccess);
+
+            std::vector<float> output(output_elems);
+            EXPECT_EQ(hipMemcpy(output.data(), obuf, obytes, hipMemcpyDeviceToHost), hipSuccess);
+            EXPECT_EQ(rocfft_plan_destroy(plan), rocfft_status_success);
+            return output;
+        };
+
+        // baseline, with nothing else in the twiddle cache
+        rocfft_cleanup();
+        rocfft_setup();
+        const std::vector<float> expected = real_inverse();
+
+        // again, but after a complex transform of the colliding lengths whose
+        // plan, and so whose cached twiddle table, is still alive
+        rocfft_cleanup();
+        rocfft_setup();
+
+        rocfft_plan complex_plan = nullptr;
+        ASSERT_EQ(rocfft_plan_create(&complex_plan,
+                                     rocfft_placement_inplace,
+                                     rocfft_transform_type_complex_forward,
+                                     rocfft_precision_single,
+                                     complex_lengths.size(),
+                                     complex_lengths.data(),
+                                     1,
+                                     nullptr),
+                  rocfft_status_success);
+
+        const size_t complex_bytes
+            = complex_lengths[0] * complex_lengths[1] * sizeof(rocfft_complex<float>);
+        gpubuf_t<rocfft_complex<float>> complex_dev;
+        ASSERT_EQ(complex_dev.alloc(complex_bytes), hipSuccess);
+        ASSERT_EQ(hipMemset(complex_dev.data(), 0, complex_bytes), hipSuccess);
+        void* complex_buf = complex_dev.data();
+        ASSERT_EQ(rocfft_execute(complex_plan, &complex_buf, nullptr, nullptr),
+                  rocfft_status_success);
+        ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
+
+        const std::vector<float> actual = real_inverse();
+
+        ASSERT_EQ(rocfft_plan_destroy(complex_plan), rocfft_status_success);
+
+        ASSERT_EQ(actual.size(), expected.size());
+        for(size_t i = 0; i < expected.size(); ++i)
+            ASSERT_FLOAT_EQ(actual[i], expected[i]) << "element " << i;
+    }
+    ROCFFT_CATCH_TEST_EXCEPTIONS
+}
+
+// check that a multi-dimensional real forward transform gets its own twiddle
+// table when a plain complex kernel of the same size already has one cached.
+// The real transform runs on a complex kernel with fused post-processing, which
+// needs a half-N section appended that the complex kernel does not use.
+TEST(rocfft_UnitTest, twiddle_cache_2D_real_forward_after_complex)
+{
+    PROB_SKIP_UNITTEST();
+    try
+    {
+        LocalCleanup restore([]() {
+            rocfft_cleanup();
+            rocfft_setup();
+        });
+
+        // a real length of 8 on one dimension and 4 on the other decomposes to
+        // a 4x4 complex kernel, so it collides with the complex transform below
+        const std::vector<size_t> real_lengths{8, 4};
+        const std::vector<size_t> complex_lengths{4, 4};
+
+        const size_t input_elems  = real_lengths[0] * real_lengths[1];
+        const size_t output_elems = (real_lengths[0] / 2 + 1) * real_lengths[1];
+
+        // any real input will do - the property under test is only that the
+        // same input gives the same output
+        std::vector<float> input(input_elems);
+        for(size_t i = 0; i < input.size(); ++i)
+            input[i] = static_cast<float>(i % 7) - 3.0f;
+
+        auto real_forward = [&]() {
+            const size_t                    ibytes = input.size() * sizeof(input[0]);
+            const size_t                    obytes = output_elems * sizeof(rocfft_complex<float>);
+            gpubuf_t<float>                 idev;
+            gpubuf_t<rocfft_complex<float>> odev;
+            EXPECT_EQ(idev.alloc(ibytes), hipSuccess);
+            EXPECT_EQ(odev.alloc(obytes), hipSuccess);
+            EXPECT_EQ(hipMemcpy(idev.data(), input.data(), ibytes, hipMemcpyHostToDevice),
+                      hipSuccess);
+
+            rocfft_plan plan = nullptr;
+            EXPECT_EQ(rocfft_plan_create(&plan,
+                                         rocfft_placement_notinplace,
+                                         rocfft_transform_type_real_forward,
+                                         rocfft_precision_single,
+                                         real_lengths.size(),
+                                         real_lengths.data(),
+                                         1,
+                                         nullptr),
+                      rocfft_status_success);
+            void* ibuf = idev.data();
+            void* obuf = odev.data();
+            EXPECT_EQ(rocfft_execute(plan, &ibuf, &obuf, nullptr), rocfft_status_success);
+            EXPECT_EQ(hipDeviceSynchronize(), hipSuccess);
+
+            std::vector<rocfft_complex<float>> output(output_elems);
+            EXPECT_EQ(hipMemcpy(output.data(), obuf, obytes, hipMemcpyDeviceToHost), hipSuccess);
+            EXPECT_EQ(rocfft_plan_destroy(plan), rocfft_status_success);
+            return output;
+        };
+
+        // baseline, with nothing else in the twiddle cache
+        rocfft_cleanup();
+        rocfft_setup();
+        const std::vector<rocfft_complex<float>> expected = real_forward();
+
+        // again, but after a complex transform of the colliding lengths whose
+        // plan, and so whose cached twiddle table, is still alive
+        rocfft_cleanup();
+        rocfft_setup();
+
+        rocfft_plan complex_plan = nullptr;
+        ASSERT_EQ(rocfft_plan_create(&complex_plan,
+                                     rocfft_placement_inplace,
+                                     rocfft_transform_type_complex_forward,
+                                     rocfft_precision_single,
+                                     complex_lengths.size(),
+                                     complex_lengths.data(),
+                                     1,
+                                     nullptr),
+                  rocfft_status_success);
+        // hold the plan for the rest of the test, however the test exits
+        LocalCleanup destroy_complex_plan(
+            [&]() { EXPECT_EQ(rocfft_plan_destroy(complex_plan), rocfft_status_success); });
+
+        const size_t complex_bytes
+            = complex_lengths[0] * complex_lengths[1] * sizeof(rocfft_complex<float>);
+        gpubuf_t<rocfft_complex<float>> complex_dev;
+        ASSERT_EQ(complex_dev.alloc(complex_bytes), hipSuccess);
+        ASSERT_EQ(hipMemset(complex_dev.data(), 0, complex_bytes), hipSuccess);
+        void* complex_buf = complex_dev.data();
+        ASSERT_EQ(rocfft_execute(complex_plan, &complex_buf, nullptr, nullptr),
+                  rocfft_status_success);
+        ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
+
+        const std::vector<rocfft_complex<float>> actual = real_forward();
+
+        ASSERT_EQ(actual.size(), expected.size());
+        for(size_t i = 0; i < expected.size(); ++i)
+        {
+            ASSERT_FLOAT_EQ(actual[i].x, expected[i].x) << "element " << i << " real part";
+            ASSERT_FLOAT_EQ(actual[i].y, expected[i].y) << "element " << i << " imaginary part";
+        }
+    }
+    ROCFFT_CATCH_TEST_EXCEPTIONS
+}
+
 // Check whether logs can be emitted from multiple threads properly
 TEST(rocfft_UnitTest, log_multithreading)
 {

projects/rocfft/library/src/include/repo.h

@@ -66,9 +66,12 @@ class Repo
     // key structure for 2D twiddles
     struct repo_twd_key_2D_t
     {
-        size_t              length0   = 0;
-        size_t              length1   = 0;
-        rocfft_precision    precision = rocfft_precision_single;
+        size_t           length0   = 0;
+        size_t           length1   = 0;
+        rocfft_precision precision = rocfft_precision_single;
+        // changes the table twiddles_create_2D generates, so it belongs in the
+        // key. attach_halfN2 does not: it is always false at the only call site.
+        bool                attach_halfN1 = false;
         std::vector<size_t> radices1;
         std::vector<size_t> radices2;
         // buffers are in device memory, so we need per-device
@@ -83,6 +86,8 @@ class Repo
                 return length1 < other.length1;
             if(precision != other.precision)
                 return precision < other.precision;
+            if(attach_halfN1 != other.attach_halfN1)
+                return attach_halfN1 < other.attach_halfN1;
             if(radices1 != other.radices1)
                 return radices1 < other.radices1;
             if(radices2 != other.radices2)

projects/rocfft/library/src/repo.cpp

@@ -204,7 +204,7 @@ std::pair<void*, size_t> Repo::GetTwiddles2D(size_t                     length0,
     std::lock_guard<std::mutex> lck(mtx);
     Repo&                       repo = Repo::GetRepo();
 
-    repo_twd_key_2D_t key{length0, length1, precision, radices1, radices2};
+    repo_twd_key_2D_t key{length0, length1, precision, attach_halfN1, radices1, radices2};
     return GetTwiddlesInternal(
         key, repo.twiddles_2D, repo.twiddles_2D_reverse, [&](unsigned int deviceId) {
             return twiddles_create_2D(length0,